Support running modules

pyinstrument/__main__.py

@@ -2,6 +2,7 @@
 import sys
 import os
 import codecs
+import runpy
 from pyinstrument import Profiler
 from pyinstrument.profiler import get_recorder_class, get_renderer_class
 from .six import exec_
@@ -44,17 +45,31 @@ def main():
         dest='color', action='store_false',
         help='force no color text output')
 
+    parser.add_option('-m', '',
+        dest='module_name', action='store',
+        help='searches sys.path for the named module and runs the  corresponding .py file as a script.')
+
     if not sys.argv[1:]:
         parser.print_help()
         sys.exit(2)
 
-    (options, args) = parser.parse_args()
+    options, args = parser.parse_args()
+
+    if args == [] and options.module_name is None:
+        parser.print_help()
+        sys.exit(2)
+
     sys.argv[:] = args
 
-    if len(args) > 0:
+    if options.module_name is not None:
+        code = "run_module(modname, run_name='__main__')"
+        globs = {
+            'run_module': runpy.run_module,
+            'modname': options.module_name
+        }
+    else:
         progname = args[0]
         sys.path.insert(0, os.path.dirname(progname))
-
         with open(progname, 'rb') as fp:
             code = compile(fp.read(), progname, 'exec')
         globs = {
@@ -63,45 +78,43 @@ def main():
             '__package__': None,
         }
 
-        if options.output_renderer:
-            renderer = options.output_renderer
-        elif options.output_html:
-            renderer = 'html'
-        else:
-            renderer = 'text'
+    if options.output_renderer:
+        renderer = options.output_renderer
+    elif options.output_html:
+        renderer = 'html'
+    else:
+        renderer = 'text'
 
-        recorder = get_renderer_class(renderer).preferred_recorder
+    recorder = get_renderer_class(renderer).preferred_recorder
 
-        profiler = Profiler(recorder=recorder)
+    profiler = Profiler(recorder=recorder)
 
-        profiler.start()
+    profiler.start()
 
-        try:
-            exec_(code, globs, None)
-        except (SystemExit, KeyboardInterrupt):
-            pass
+    try:
+        exec_(code, globs, None)
+    except (SystemExit, KeyboardInterrupt):
+        pass
 
-        profiler.stop()
+    profiler.stop()
 
-        if options.outfile:
-            f = codecs.open(options.outfile, 'w', 'utf-8')
-        else:
-            f = sys.stdout
+    if options.outfile:
+        f = codecs.open(options.outfile, 'w', 'utf-8')
+    else:
+        f = sys.stdout
 
-        renderer_kwargs = {}
+    renderer_kwargs = {}
 
-        if renderer == 'text':
-            unicode_override = options.unicode != None
-            color_override = options.color != None
-            unicode = options.unicode if unicode_override else file_supports_unicode(f)
-            color = options.color if color_override else file_supports_color(f)
-            
-            renderer_kwargs = {'unicode': unicode, 'color': color}
+    if renderer == 'text':
+        unicode_override = options.unicode != None
+        color_override = options.color != None
+        unicode = options.unicode if unicode_override else file_supports_unicode(f)
+        color = options.color if color_override else file_supports_color(f)
+
+        renderer_kwargs = {'unicode': unicode, 'color': color}
 
-        f.write(profiler.output(renderer=renderer, **renderer_kwargs))
-        f.close()
-    else:
-        parser.print_usage()
+    f.write(profiler.output(renderer=renderer, **renderer_kwargs))
+    f.close()
 
 
 def file_supports_color(file_obj):

test/test_cmdline.py

@@ -9,11 +9,10 @@ def test_command_line():
     assert 'busy_wait' in str(output)
     assert 'do_nothing' in str(output)
 
-@pytest.mark.xfail
 def test_module_running():
-    working_dir = os.path.dirname(__file__)
+    working_dir = os.path.normpath(os.path.join(os.path.dirname(__file__), '..'))
     output = subprocess.check_output(
-        [sys.executable, '-m', 'pyinstrument', '-m', 'busywait_module'],
+        [sys.executable, '-m', 'pyinstrument', '-m', 'test.busywait_module'],
         cwd=working_dir
     )