Merge pull request #271 from joerick/dash-c-program

Add '-c' feature, so code can be passed on the command line

pyinstrument/__main__.py

@@ -34,13 +34,25 @@ def main():
     parser = optparse.OptionParser(usage=usage, version=version_string)
     parser.allow_interspersed_args = False
 
-    def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionParser):
-        parser.values.module_name = value  # type: ignore
-
-        # everything after the -m argument should be passed to that module
-        parser.values.module_args = parser.rargs + parser.largs  # type: ignore
-        parser.rargs[:] = []  # type: ignore
-        parser.largs[:] = []  # type: ignore
+    def store_and_consume_remaining(
+        option: optparse.Option, opt: str, value: str, parser: optparse.OptionParser
+    ):
+        """
+        A callback for optparse that stores the value and consumes all
+        remaining arguments, storing them in the same variable as a tuple.
+        """
+
+        # assert a few things we know to be true about the parser
+        assert option.dest
+        assert parser.rargs is not None
+        assert parser.largs is not None
+
+        # everything after this argument should be consumed
+        remaining_arguments = parser.rargs + parser.largs
+        parser.rargs[:] = []
+        parser.largs[:] = []
+
+        setattr(parser.values, option.dest, ValueWithRemainingArgs(value, remaining_arguments))
 
     parser.add_option(
         "--load",
@@ -62,12 +74,21 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
     parser.add_option(
         "-m",
         "",
-        dest="module_name",
+        dest="module",
         action="callback",
-        callback=dash_m_callback,
-        type="str",
+        callback=store_and_consume_remaining,
+        type="string",
         help="run library module as a script, like 'python -m module'",
     )
+    parser.add_option(
+        "-c",
+        "",
+        dest="program",
+        action="callback",
+        callback=store_and_consume_remaining,
+        type="string",
+        help="program passed in as string, like 'python -c \"...\"'",
+    )
     parser.add_option(
         "",
         "--from-path",
@@ -244,7 +265,8 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
     session_options_used = [
         options.load is not None,
         options.load_prev is not None,
-        options.module_name is not None,
+        options.module is not None,
+        options.program is not None,
         len(args) > 0,
     ]
     if session_options_used.count(True) == 0:
@@ -253,7 +275,7 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
     if session_options_used.count(True) > 1:
         parser.error("You can only specify one of --load, --load-prev, -m, or script arguments")
 
-    if options.module_name is not None and options.from_path:
+    if options.module is not None and options.from_path:
         parser.error("The options -m and --from-path are mutually exclusive.")
 
     if options.from_path and sys.platform == "win32":
@@ -297,14 +319,21 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
     elif options.load:
         session = Session.load(options.load)
     else:
-        if options.module_name is not None:
+        # we are running some code
+        if options.module is not None:
             if not (sys.path[0] and os.path.samefile(sys.path[0], ".")):
                 # when called with '-m', search the cwd for that module
                 sys.path[0] = os.path.abspath(".")
 
-            argv = [options.module_name] + options.module_args
+            argv = [options.module.value] + options.module.remaining_args
             code = "run_module(modname, run_name='__main__', alter_sys=True)"
-            globs = {"run_module": runpy.run_module, "modname": options.module_name}
+            globs = {"run_module": runpy.run_module, "modname": options.module.value}
+        elif options.program is not None:
+            argv = ["-c", *options.program.remaining_args]
+            code = options.program.value
+            globs = {"__name__": "__main__"}
+            # set the first path entry to '' to match behaviour of python -c
+            sys.path[0] = ""
         else:
             argv = args
             if options.from_path:
@@ -322,15 +351,15 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
             code = "run_path(progname, run_name='__main__')"
             globs = {"run_path": runpy.run_path, "progname": progname}
 
+        old_argv = sys.argv.copy()
+
         # there is no point using async mode for command line invocation,
         # because it will always be capturing the whole program, we never want
         # any execution to be <out-of-context>, and it avoids duplicate
         # profiler errors.
         profiler = Profiler(interval=options.interval, async_mode="disabled")
-
         profiler.start()
 
-        old_argv = sys.argv.copy()
         try:
             sys.argv[:] = argv
             exec(code, globs, None)
@@ -552,8 +581,8 @@ class CommandLineOptions:
     A type that codifies the `options` variable.
     """
 
-    module_name: str | None
-    module_args: list[str]
+    module: ValueWithRemainingArgs | None
+    program: ValueWithRemainingArgs | None
     load: str | None
     load_prev: str | None
     from_path: str | None
@@ -573,5 +602,11 @@ class CommandLineOptions:
     interval: float
 
 
+class ValueWithRemainingArgs:
+    def __init__(self, value: str, remaining_args: list[str]):
+        self.value = value
+        self.remaining_args = remaining_args
+
+
 if __name__ == "__main__":
     main()

test/test_cmdline.py

@@ -2,14 +2,13 @@
 import re
 import subprocess
 import sys
+import textwrap
 from pathlib import Path
 
 import pytest
 
 from .util import BUSY_WAIT_SCRIPT
 
-# this script just does a busywait for 0.25 seconds.
-
 EXECUTION_DETAILS_SCRIPT = f"""
 #!{sys.executable}
 import sys, os
@@ -78,6 +77,35 @@ def test_path(self, pyinstrument_invocation, tmp_path: Path, monkeypatch):
             [*pyinstrument_invocation, "--from-path", "--", "pyi_test_program"],
         )
 
+    def test_program_passed_as_string(self, pyinstrument_invocation, tmp_path: Path):
+        # check the program actually runs
+        output_file = tmp_path / "output.txt"
+        output = subprocess.check_output(
+            [
+                *pyinstrument_invocation,
+                "-c",
+                textwrap.dedent(
+                    f"""
+                    import sys
+                    from pathlib import Path
+                    output_file = Path(sys.argv[1])
+                    output_file.write_text("Hello World")
+                    print("Finished.")
+                    """
+                ),
+                str(output_file),
+            ],
+        )
+
+        assert "Finished." in str(output)
+        assert output_file.read_text() == "Hello World"
+
+        # check the output
+        output = subprocess.check_output([*pyinstrument_invocation, "-c", BUSY_WAIT_SCRIPT])
+
+        assert "busy_wait" in str(output)
+        assert "do_nothing" in str(output)
+
     def test_script_execution_details(self, pyinstrument_invocation, tmp_path: Path):
         program_path = tmp_path / "program.py"
         program_path.write_text(EXECUTION_DETAILS_SCRIPT)
@@ -157,6 +185,27 @@ def test_path_execution_details(self, pyinstrument_invocation, tmp_path: Path, m
         print("process_native.stderr", process_native.stderr)
         assert process_pyi.stderr == process_native.stderr
 
+    def test_program_passed_as_string_execution_details(
+        self, pyinstrument_invocation, tmp_path: Path
+    ):
+        process_pyi = subprocess.run(
+            [*pyinstrument_invocation, "-c", EXECUTION_DETAILS_SCRIPT],
+            stderr=subprocess.PIPE,
+            check=True,
+            text=True,
+        )
+        process_native = subprocess.run(
+            [sys.executable, "-c", EXECUTION_DETAILS_SCRIPT],
+            stderr=subprocess.PIPE,
+            check=True,
+            text=True,
+        )
+
+        print("process_pyi.stderr", process_pyi.stderr)
+        print("process_native.stderr", process_native.stderr)
+        assert process_native.stderr
+        assert process_pyi.stderr == process_native.stderr
+
     def test_session_save_and_load(self, pyinstrument_invocation, tmp_path: Path):
         busy_wait_py = tmp_path / "busy_wait.py"
         busy_wait_py.write_text(BUSY_WAIT_SCRIPT)