Merge pull request #259 from BarrensZeppelin/master

Add `-m` switch to vmprof cli

Author: mattip

docs/vmprof.rst

@@ -84,6 +84,8 @@ After ``-m vmprof`` you can specify some options:
 * ``-n`` - enable all C frames, only useful if you have a debug build of
   PyPy or CPython.
 
+* ``-m`` - indicate that the provided program should be executed as a module (like ``python -m``). Example: ``-m vmprof -m http.server``
+
 * ``--lines`` - enable line profiling mode. This mode adds some overhead to profiling, but in addition to function calls it marks the execution of the specific lines inside functions.
 
 * ``-o file`` - save logs for later

vmprof/__main__.py

@@ -56,9 +56,12 @@ def main():
         _jitlog.enable(fd)
     # invoke the user program:
     try:
-        sys.argv = [args.program] + args.args
-        sys.path.insert(0, os.path.dirname(args.program))
-        runpy.run_path(args.program, run_name='__main__')
+        sys.argv[1:] = args.args
+        if args.module:
+            runpy.run_module(args.program, run_name='__main__', alter_sys=True)
+        else:
+            sys.path.insert(0, os.path.dirname(args.program))
+            runpy.run_path(args.program, run_name='__main__')
     except BaseException as e:
         if not isinstance(e, (KeyboardInterrupt, SystemExit)):
             raise

vmprof/cli.py

@@ -10,6 +10,12 @@ def build_argparser():
         description='VMprof',
         prog="vmprof"
     )
+    parser.add_argument(
+        '-m',
+        dest='module',
+        action='store_true',
+        help='Run module as a script'
+    )
     parser.add_argument(
         'program',
         help='program'

vmprof/test/test_config.py

@@ -61,3 +61,10 @@ def test_parser_without_section():
 
     assert test_file == args.config.name
     assert args.no_native == True
+
+
+def test_parser_with_m_switch():
+    args = cli.parse_args(['-m', 'example'])
+
+    assert args.module
+    assert args.program == 'example'