sys.argv shouldn't be empty in the running module

joerick · Aug 4, 2018 · 2c05990 · 2c05990
1 parent 6448ccb
commit 2c05990
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/pyinstrument/__main__.py b/pyinstrument/__main__.py
@@ -59,15 +59,15 @@ def main():
         parser.print_help()
         sys.exit(2)
 
-    sys.argv[:] = args
-
     if options.module_name is not None:
+        sys.argv[:] = [options.module_name] + args
         code = "run_module(modname, run_name='__main__')"
         globs = {
             'run_module': runpy.run_module,
             'modname': options.module_name
         }
     else:
+        sys.argv[:] = args
         progname = args[0]
         sys.path.insert(0, os.path.dirname(progname))
         with open(progname, 'rb') as fp:

diff --git a/test/test_cmdline.py b/test/test_cmdline.py
@@ -10,7 +10,7 @@ def test_command_line():
     assert 'do_nothing' in str(output)
 
 def test_module_running():
-    working_dir = os.path.normpath(os.path.join(os.path.dirname(__file__), '..'))
+    working_dir = os.path.dirname(os.path.dirname(__file__))
     output = subprocess.check_output(
         [sys.executable, '-m', 'pyinstrument', '-m', 'test.busywait_module'],
         cwd=working_dir
@@ -19,3 +19,10 @@ def test_module_running():
     assert 'busy_wait' in str(output)
     assert 'do_nothing' in str(output)
 
+def test_running_yourself_as_module():
+    working_dir = os.path.dirname(os.path.dirname(__file__))
+    subprocess.check_call(
+        [sys.executable, '-m', 'pyinstrument', '-m', 'pyinstrument'],
+        cwd=working_dir
+    )
+