diff --git a/src/memray/commands/run.py b/src/memray/commands/run.py index a1a9025436..6e2ed1206a 100644 --- a/src/memray/commands/run.py +++ b/src/memray/commands/run.py @@ -51,6 +51,8 @@ def _run_tracker( try: if args.run_as_module: runpy.run_module(args.script, run_name="__main__", alter_sys=True) + elif args.run_as_cmd: + exec(args.script, {"__name__": "__main__"}) else: runpy.run_path(args.script, run_name="__main__") finally: @@ -62,6 +64,7 @@ def _child_process( port: int, native: bool, run_as_module: bool, + run_as_cmd: bool, quiet: bool, script: str, script_args: List[str], @@ -69,6 +72,7 @@ def _child_process( args = argparse.Namespace( native=native, run_as_module=run_as_module, + run_as_cmd=run_as_cmd, quiet=quiet, script=script, script_args=script_args, @@ -84,7 +88,7 @@ def _run_child_process_and_attach(args: argparse.Namespace) -> None: raise MemrayCommandError(f"Invalid port: {port}", exit_code=1) arguments = ( - f"{port},{args.native},{args.run_as_module},{args.quiet}," + f"{port},{args.native},{args.run_as_module},{args.run_as_cmd},{args.quiet}," f'"{args.script}",{args.script_args}' ) tracked_app_cmd = [ @@ -128,8 +132,12 @@ def _run_with_socket_output(args: argparse.Namespace) -> None: def _run_with_file_output(args: argparse.Namespace) -> None: if args.output is None: - output = f"memray-{os.path.basename(args.script)}.{os.getpid()}.bin" - filename = os.path.join(os.path.dirname(args.script), output) + script_name = args.script + if args.run_as_cmd: + script_name = "string" + + output = f"memray-{os.path.basename(script_name)}.{os.getpid()}.bin" + filename = os.path.join(os.path.dirname(script_name), output) else: filename = args.output @@ -163,7 +171,7 @@ class RunCommand: """Run the specified application and track memory usage""" def prepare_parser(self, parser: argparse.ArgumentParser) -> None: - parser.usage = "%(prog)s [-m module | file] [args]" + parser.usage = "%(prog)s [-m module | -c cmd | file] [args]" output_group = parser.add_mutually_exclusive_group() output_group.add_argument( "-o", @@ -218,6 +226,13 @@ def prepare_parser(self, parser: argparse.ArgumentParser) -> None: action="store_true", default=False, ) + parser.add_argument( + "-c", + help="Program passed in as string", + action="store_true", + dest="run_as_cmd", + default=False, + ) parser.add_argument( "-m", help="Run library module as a script (terminates option list)", @@ -237,11 +252,15 @@ def validate_target_file(self, args: argparse.Namespace) -> None: if args.run_as_module: return try: - source = pathlib.Path(args.script).read_bytes() + if args.run_as_cmd: + source = bytes(args.script, "UTF-8") + else: + source = pathlib.Path(args.script).read_bytes() ast.parse(source) except (SyntaxError, ValueError): raise MemrayCommandError( - "Only Python files can be executed under memray", exit_code=1 + "Only valid Python files or commands can be executed under memray", + exit_code=1, ) def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None: @@ -249,6 +268,8 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None parser.error("The --live-port argument requires --live-remote") if args.follow_fork is True and (args.live_mode or args.live_remote_mode): parser.error("--follow-fork cannot be used with the live TUI") + if args.run_as_cmd and pathlib.Path(args.script).exists(): + parser.error("remove the option -c to run a file") self.validate_target_file(args) diff --git a/tests/integration/test_main.py b/tests/integration/test_main.py index 209a604c19..0bed1ddaa9 100644 --- a/tests/integration/test_main.py +++ b/tests/integration/test_main.py @@ -254,7 +254,10 @@ def test_run_file_that_is_not_python(self, capsys, option): # THEN captured = capsys.readouterr() - assert captured.err.strip() == "Only Python files can be executed under memray" + assert ( + captured.err.strip() + == "Only valid Python files or commands can be executed under memray" + ) @patch("memray.commands.run.os.getpid") def test_run_file_exists(self, getpid, tmp_path, monkeypatch, capsys): diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index 4df4c16a4b..1dff5e89ca 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -36,7 +36,7 @@ def test_run_without_arguments( main(["run"]) captured = capsys.readouterr() - assert "usage: memray run [-m module | file] [args]" in captured.err + assert "usage: memray run [-m module | -c cmd | file] [args]" in captured.err def test_run_default_output( self, getpid_mock, runpy_mock, tracker_mock, validate_mock @@ -94,6 +94,22 @@ def test_run_module(self, getpid_mock, runpy_mock, tracker_mock, validate_mock): "foobar", run_name="__main__", alter_sys=True ) + def test_run_cmd_is_validated( + self, getpid_mock, runpy_mock, tracker_mock, validate_mock + ): + with patch.object(RunCommand, "validate_target_file"): + assert 0 == main(["run", "-c", "x = [i for i in range(10)]"]) + with pytest.raises(SyntaxError): + main(["run", "-c", "[i for i in range(10)"]) + + def test_run_cmd(self, getpid_mock, runpy_mock, tracker_mock, validate_mock): + with patch("memray.commands.run.exec") as mock_exec: + assert 0 == main(["run", "-c", "x = 10; y = abs(-10)"]) + assert not runpy_mock.called + mock_exec.assert_called_with( + "x = 10; y = abs(-10)", {"__name__": "__main__"} + ) + def test_run_file(self, getpid_mock, runpy_mock, tracker_mock, validate_mock): with patch.object(RunCommand, "validate_target_file"): assert 0 == main(["run", "foobar.py", "arg1", "arg2"]) @@ -136,7 +152,7 @@ def test_run_with_live( sys.executable, "-c", "from memray.commands.run import _child_process;" - '_child_process(1234,False,False,False,"./directory/foobar.py",' + '_child_process(1234,False,False,False,False,"./directory/foobar.py",' "['arg1', 'arg2'])", ], stderr=-1,