Add -c option: program passed as string

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,13 +64,15 @@ def _child_process(
     port: int,
     native: bool,
     run_as_module: bool,
+    run_as_cmd: bool,
     quiet: bool,
     script: str,
     script_args: List[str],
 ) -> None:
     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,18 +252,24 @@ 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:
         if args.live_port is not None and not args.live_remote_mode:
             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)
 

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):

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,