From 296d7d172579a14cc542c2ab96428b2ee75ea61b Mon Sep 17 00:00:00 2001
From: Rasmus Wriedt Larsen <rasmuswl@github.com>
Date: Tue, 21 Jul 2020 19:39:51 +0200
Subject: [PATCH] Python: CG trace: Allow tracing modules

As would normally be invoked by `python -m <module-name>` now works with
`cg-trace --module <module-name>`.

This is useful for tracing invocations of `pytest`.
---
 .../src/cg_trace/cmdline.py                   |  4 ++
 .../src/cg_trace/main.py                      | 37 +++++++++++++------
 2 files changed, 30 insertions(+), 11 deletions(-)

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
index ddbeeaf50dc..d17328638dc 100644
--- 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
index 3691fe4149b..42eb361c6db 100644
--- 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)
+    if opts.module:
+        import runpy
 
-    with open(opts.progname) as fp:
-        code = compile(fp.read(), opts.progname, "exec")
+        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)
 
-    # try to emulate __main__ namespace as much as possible
-    globs = {
-        "__file__": opts.progname,
-        "__name__": "__main__",
-        "__package__": None,
-        "__cached__": None,
-    }
+        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(