Repository: bloomberg/memray (python)

Issue:
`memray run` overwrites `sys.argv[0]` even when `-I` or `-P` is used
`memray run` tries to match the updates to `sys.path` that are performed by the interpreter itself when it starts. When those path updates are suppressed (by the `-I` or `-P` command line flags or their corresponding environment variables), Memray's emulation winds up overwriting paths in `sys.path` that needed to be left alone.

Pull request description:
Fix `sys.path` manipulations for isolated mode
Closes #551 

I struggled to write good tests for this, and while working on it I realized that at least one of the existing tests wasn't actually proving what it set out to prove, so I restructured the existing tests as well.

Patch:
diff --git a/news/552.bugfix.rst b/news/552.bugfix.rst
new file mode 100644
index 0000000..939ecb8
--- /dev/null
+++ b/news/552.bugfix.rst
@@ -0,0 +1,1 @@
+Fix a bug causing the first entry of ``sys.path`` to be erroneously overwritten by ``memray run`` when the Python interpreter was launched with the ``-I`` or ``-P`` flag, or when the ``PYTHONSAFEPATH`` environment variable was set.
diff --git a/src/memray/commands/run.py b/src/memray/commands/run.py
index 2808ade..4f006cf 100644
--- a/src/memray/commands/run.py
+++ b/src/memray/commands/run.py
@@ -31,6 +31,12 @@ def _get_free_port() -> int:
         return int(sock.getsockname()[1])
 
 
+def _should_modify_sys_path() -> bool:
+    isolated_mode = sys.flags.isolated
+    safe_path_mode = getattr(sys.flags, "safe_path", False)  # New in Python 3.11
+    return not isolated_mode and not safe_path_mode
+
+
 def _run_tracker(
     destination: Destination,
     args: argparse.Namespace,
@@ -52,16 +58,21 @@ def _run_tracker(
         pid = os.getpid()
         try:
             if args.run_as_module:
-                sys.path[0] = os.getcwd()
+                if _should_modify_sys_path():
+                    sys.path[0] = os.getcwd()
                 # run_module will replace argv[0] with the script's path
                 sys.argv = ["", *args.script_args]
                 runpy.run_module(args.script, run_name="__main__", alter_sys=True)
             elif args.run_as_cmd:
-                sys.path[0] = ""
+                if _should_modify_sys_path():
+                    sys.path[0] = ""
                 sys.argv = ["-c", *args.script_args]
                 exec(args.script, {"__name__": "__main__"})
             else:
-                sys.path[0] = str(pathlib.Path(args.script).resolve().parent.absolute())
+                if _should_modify_sys_path():
+                    sys.path[0] = str(
+                        pathlib.Path(args.script).resolve().parent.absolute()
+                    )
                 sys.argv = [args.script, *args.script_args]
                 runpy.run_path(args.script, run_name="__main__")
         finally:


Test patch:
diff --git a/tests/integration/test_main.py b/tests/integration/test_main.py
index ee9d3aa..d724a1f 100644
--- a/tests/integration/test_main.py
+++ b/tests/integration/test_main.py
@@ -268,9 +268,7 @@ class TestRunSubcommand:
         # GIVEN
         out_file = tmp_path / "result.bin"
         target_file = tmp_path / "test.py"
-        target_file.write_text("import some_adjacent_module")
-        other_file = tmp_path / "some_adjacent_module.py"
-        other_file.write_text("import sys; print(sys.argv); print(sys.path[0])")
+        target_file.write_text("import json, sys; print(json.dumps(sys.path))")
 
         # WHEN
         proc = subprocess.run(
@@ -292,57 +290,127 @@ class TestRunSubcommand:
         )
 
         # THEN
-        assert proc.returncode == 0
-        argv, path0 = proc.stdout.splitlines()
-        assert argv == repr([str(target_file), "some", "provided args"])
-        assert path0 == str(tmp_path)
         assert out_file.exists()
+        assert proc.returncode == 0
+        # Running `python -m` put cwd in sys.path; ensure we replaced it.
+        path = json.loads(proc.stdout)
+        assert os.getcwd() not in path
+        assert str(tmp_path) in path
 
-    def test_sys_manipulations_when_running_module(self, tmp_path):
+    @pytest.mark.parametrize(
+        "isolation_flag", ["-I"] + (["-P"] if sys.version_info > (3, 11) else [])
+    )
+    def test_suppressing_sys_manipulations_when_running_script(
+        self, tmp_path, isolation_flag
+    ):
         # GIVEN
         out_file = tmp_path / "result.bin"
         target_file = tmp_path / "test.py"
-        target_file.write_text("import some_adjacent_module")
-        other_file = tmp_path / "some_adjacent_module.py"
-        other_file.write_text("import sys; print(sys.argv); print(sys.path[0])")
-        env = os.environ.copy()
-        env["PYTHONPATH"] = str(tmp_path) + ":" + os.environ.get("PYTHONPATH", "")
+        target_file.write_text("import json, sys; print(json.dumps(sys.path))")
 
         # WHEN
         proc = subprocess.run(
             [
                 sys.executable,
+                isolation_flag,
                 "-m",
                 "memray",
                 "run",
                 "--quiet",
                 "--output",
                 str(out_file),
-                "-m",
-                "test",
+                str(target_file),
                 "some",
                 "provided args",
             ],
             check=True,
             capture_output=True,
             text=True,
-            env=env,
         )
 
         # THEN
+        assert out_file.exists()
         assert proc.returncode == 0
-        argv, path0 = proc.stdout.splitlines()
-        assert argv == repr([str(target_file), "some", "provided args"])
-        assert path0 == os.getcwd()
+        # Running `python -m` did not put cwd in sys.path; ensure it isn't
+        # there, and neither is the tmp_path we would have replaced it with.
+        path = json.loads(proc.stdout)
+        assert os.getcwd() not in path
+        assert str(tmp_path) not in path
+
+    def test_sys_manipulations_when_running_module(self, tmp_path):
+        # GIVEN
+        out_file = tmp_path / "result.bin"
+
+        # WHEN
+        proc = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                "import sys; from memray.commands import main; sys.exit(main())",
+                "run",
+                "--quiet",
+                "--output",
+                str(out_file),
+                "-m",
+                "site",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+        # THEN
         assert out_file.exists()
+        assert proc.returncode == 0
+        # Running `python -c` put "" in sys.path; ensure we replaced it.
+        path = eval(
+            " ".join(line for line in proc.stdout.splitlines() if line.startswith(" "))
+        )
+        assert "" not in path
+        assert os.getcwd() in path
+
+    @pytest.mark.parametrize(
+        "isolation_flag", ["-I"] + (["-P"] if sys.version_info > (3, 11) else [])
+    )
+    def test_suppressing_sys_manipulations_when_running_module(
+        self, tmp_path, isolation_flag
+    ):
+        # GIVEN
+        out_file = tmp_path / "result.bin"
+
+        # WHEN
+        proc = subprocess.run(
+            [
+                sys.executable,
+                isolation_flag,
+                "-c",
+                "import sys; from memray.commands import main; sys.exit(main())",
+                "run",
+                "--quiet",
+                "--output",
+                str(out_file),
+                "-m",
+                "site",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+        # THEN
+        assert out_file.exists()
+        assert proc.returncode == 0
+        # Running `python -c` did not put "" in sys.path; ensure it isn't
+        # there, and neither is the os.getcwd() we would have replaced it with.
+        path = eval(
+            " ".join(line for line in proc.stdout.splitlines() if line.startswith(" "))
+        )
+        assert "" not in path
+        assert os.getcwd() not in path
 
     def test_sys_manipulations_when_running_cmd(self, tmp_path):
         # GIVEN
         out_file = tmp_path / "result.bin"
-        other_file = tmp_path / "some_adjacent_module.py"
-        other_file.write_text("import sys; print(sys.argv); print(sys.path[0])")
-        env = os.environ.copy()
-        env["PYTHONPATH"] = str(tmp_path) + ":" + os.environ.get("PYTHONPATH", "")
 
         # WHEN
         proc = subprocess.run(
@@ -355,22 +423,57 @@ class TestRunSubcommand:
                 "--output",
                 str(out_file),
                 "-c",
-                "import some_adjacent_module",
-                "some",
-                "provided args",
+                "import json, sys; print(json.dumps(sys.path))",
             ],
             check=True,
             capture_output=True,
             text=True,
-            env=env,
         )
 
         # THEN
+        assert out_file.exists()
         assert proc.returncode == 0
-        argv, path0 = proc.stdout.splitlines()
-        assert argv == repr(["-c", "some", "provided args"])
-        assert path0 == ""
+        # Running `python -m` put cwd in sys.path; ensure we replaced it.
+        path = json.loads(proc.stdout)
+        assert os.getcwd() not in path
+        assert "" in path
+
+    @pytest.mark.parametrize(
+        "isolation_flag", ["-I"] + (["-P"] if sys.version_info > (3, 11) else [])
+    )
+    def test_suppressing_sys_manipulations_when_running_cmd(
+        self, tmp_path, isolation_flag
+    ):
+        # GIVEN
+        out_file = tmp_path / "result.bin"
+
+        # WHEN
+        proc = subprocess.run(
+            [
+                sys.executable,
+                isolation_flag,
+                "-m",
+                "memray",
+                "run",
+                "--quiet",
+                "--output",
+                str(out_file),
+                "-c",
+                "import json, sys; print(json.dumps(sys.path))",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+        # THEN
         assert out_file.exists()
+        assert proc.returncode == 0
+        # Running `python -m` did not put cwd in sys.path; ensure it isn't
+        # there, and neither is the "" we would have replaced it with.
+        path = json.loads(proc.stdout)
+        assert os.getcwd() not in path
+        assert "" not in path
 
     @pytest.mark.parametrize("option", [None, "--live", "--live-remote"])
     def test_run_file_that_is_not_python(self, capsys, option):