From ce2989db56db57cf4a108efc7fa10e0a59d4fc42 Mon Sep 17 00:00:00 2001 From: Joe Rickerby Date: Fri, 6 Oct 2023 16:30:37 +0100 Subject: [PATCH 1/2] Add '-c' feature, so code can be passed on the command line --- pyinstrument/__main__.py | 73 +++++++++++++++++++++++++++++----------- test/test_cmdline.py | 50 +++++++++++++++++++++++++-- 2 files changed, 102 insertions(+), 21 deletions(-) diff --git a/pyinstrument/__main__.py b/pyinstrument/__main__.py index 4a60fd05..e63df3df 100644 --- a/pyinstrument/__main__.py +++ b/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 , 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() diff --git a/test/test_cmdline.py b/test/test_cmdline.py index 3a0706d1..bf51bae7 100644 --- a/test/test_cmdline.py +++ b/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,32 @@ 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 = subprocess.check_output( + [ + *pyinstrument_invocation, + "-c", + textwrap.dedent( + f""" + from pathlib import Path + output_file = Path("{tmp_path}/output.txt") + output_file.write_text("Hello World") + print("Finished.") + """ + ), + ], + ) + + assert "Finished." in str(output) + assert (tmp_path / "output.txt").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 +182,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) From b227322f8d9407c895be4185bd6235185bee6bca Mon Sep 17 00:00:00 2001 From: Joe Rickerby Date: Fri, 6 Oct 2023 17:40:55 +0100 Subject: [PATCH 2/2] Fix bug on windows with backslashes in string interpolation --- test/test_cmdline.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/test/test_cmdline.py b/test/test_cmdline.py index bf51bae7..d718a2c8 100644 --- a/test/test_cmdline.py +++ b/test/test_cmdline.py @@ -79,23 +79,26 @@ def test_path(self, pyinstrument_invocation, tmp_path: Path, monkeypatch): 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("{tmp_path}/output.txt") + output_file = Path(sys.argv[1]) output_file.write_text("Hello World") print("Finished.") """ ), + str(output_file), ], ) assert "Finished." in str(output) - assert (tmp_path / "output.txt").read_text() == "Hello World" + assert output_file.read_text() == "Hello World" # check the output output = subprocess.check_output([*pyinstrument_invocation, "-c", BUSY_WAIT_SCRIPT])