Address review

Author: NoahStapp

CONTRIBUTING.md

@@ -472,5 +472,5 @@ To profile a test script and generate a flame graph, follow these steps:
   ```bash
   just flamegraph <output_name=profile> <sample_rate=2000> <path/to/script>
   ```
-4. Profiling should be done on a Linux system, as macOS and Windows do not support the `--native` option of `py-spy`.
+4. If you need to include native code (for example the C extensions), profiling should be done on a Linux system, as macOS and Windows do not support the `--native` option of `py-spy`.
    Creating an ubuntu Evergreen spawn host and using `scp` to copy the flamegraph `.svg` file back to your local machine is the best way to do this.

tools/generate_flamegraph.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate a flamegraph of a given Python script.")
+
+    parser.add_argument(
+        "--output",
+        default="profile",
+        help="Output filename (default: 'profile')",
+    )
+    parser.add_argument(
+        "--sampling_rate",
+        default="2000",
+        help="Sampling rate in samples/sec (default: 2000)",
+    )
+    parser.add_argument(
+        "--native",
+        default=False,
+        action=argparse.BooleanOptionalAction,
+        help="Whether to profile native extensions (default: False)",
+    )
+    parser.add_argument(
+        "--script_path",
+        required=True,
+        help="Path to the Python script to be profiled (required)",
+    )
+
+    args = parser.parse_args()
+
+    bash_command = [
+        "py-spy",
+        "record",
+        "-o",
+        f"{args.output}.svg",
+        "-r",
+        f"{args.sampling_rate}",
+        "--",
+        "python",
+        f"{args.script_path}",
+    ]
+
+    if args.native:
+        # Insert --native option at the correct position
+        bash_command.insert(6, "--native")
+
+    try:
+        subprocess.check_call(bash_command)  # noqa: S603
+    except Exception:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()