Python: CG trace: Allow tracing modules

RasmusWL · RasmusWL · commit 296d7d172579 · 2020-07-21T19:39:51.000+02:00
As would normally be invoked by `python -m &lt;module-name&gt;` now works with
`cg-trace --module &lt;module-name&gt;`.

This is useful for tracing invocations of `pytest`.
diff --git a/python/tools/recorded-call-graph-metrics/src/cg_trace/cmdline.py b/python/tools/recorded-call-graph-metrics/src/cg_trace/cmdline.py
@@ -6,6 +6,10 @@ def parse(args):
 
     parser.add_argument("--xml")
 
+    parser.add_argument(
+        "--module", action="store_true", default=False, help="Trace a module"
+    )
+
     parser.add_argument("progname", help="file to run as main program")
     parser.add_argument(
         "arguments", nargs=argparse.REMAINDER, help="arguments to the program"
diff --git a/python/tools/recorded-call-graph-metrics/src/cg_trace/main.py b/python/tools/recorded-call-graph-metrics/src/cg_trace/main.py
@@ -41,19 +41,34 @@ def main(args=None) -> int:
 
     # These details of setting up the program to be run is very much inspired by `trace`
     # from the standard library
-    sys.argv = [opts.progname, *opts.arguments]
-    sys.path[0] = os.path.dirname(opts.progname)
-
-    with open(opts.progname) as fp:
-        code = compile(fp.read(), opts.progname, "exec")
-
-    # try to emulate __main__ namespace as much as possible
-    globs = {
-        "__file__": opts.progname,
-        "__name__": "__main__",
-        "__package__": None,
-        "__cached__": None,
-    }
+    if opts.module:
+        import runpy
+
+        module_name = opts.progname
+        _mod_name, mod_spec, code = runpy._get_module_details(module_name)
+        sys.argv = [code.co_filename, *opts.arguments]
+        globs = {
+            "__name__": "__main__",
+            "__file__": code.co_filename,
+            "__package__": mod_spec.parent,
+            "__loader__": mod_spec.loader,
+            "__spec__": mod_spec,
+            "__cached__": None,
+        }
+    else:
+        sys.argv = [opts.progname, *opts.arguments]
+        sys.path[0] = os.path.dirname(opts.progname)
+
+        with open(opts.progname) as fp:
+            code = compile(fp.read(), opts.progname, "exec")
+
+        # try to emulate __main__ namespace as much as possible
+        globs = {
+            "__file__": opts.progname,
+            "__name__": "__main__",
+            "__package__": None,
+            "__cached__": None,
+        }
 
     start = time.time()
     recorded_calls, captured_stdout, captured_stderr, exit_status = record_calls(
