Skip to content

Commit 6c4dd14

Browse files
committed
gh-158540: Add the profiled script's directory to sys.path
When a script is profiled with ``python -m profiling.sampling run script.py`` from another directory, the script cannot import modules placed next to it, because ``_sync_coordinator._execute_script()`` executes it with the working directory (added by ``_setup_environment()`` for the module case) as ``sys.path[0]`` instead of the script's own directory. Make the script's directory importable in ``_execute_script()``, matching the behavior of ``python script.py``. Add a regression test that runs the coordinator on a script importing a sibling module.
1 parent 763b6ed commit 6c4dd14

3 files changed

Lines changed: 97 additions & 0 deletions

File tree

‎Lib/profiling/sampling/_sync_coordinator.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,15 @@ def _execute_script(script_path: str, script_args: List[str], cwd: str) -> None:
168168
if not os.path.isfile(script_path):
169169
raise TargetError(f"Script not found: {script_path}")
170170

171+
# gh-158540: ``python script.py`` puts the directory containing the
172+
# script (not the current working directory) on sys.path so that modules
173+
# next to the script can be imported. _setup_environment() added the
174+
# working directory for module (``-m``) imports; make the script's
175+
# directory importable as well.
176+
script_dir = os.path.dirname(script_path)
177+
if script_dir not in sys.path:
178+
sys.path.insert(0, script_dir)
179+
171180
# Replace sys.argv to match original script call
172181
sys.argv = [script_path] + script_args
173182

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
"""Tests for the sampling profiler's synchronization coordinator."""
2+
3+
import os
4+
import socket
5+
import subprocess
6+
import sys
7+
import tempfile
8+
import unittest
9+
10+
from test.support import SHORT_TIMEOUT, requires_subprocess
11+
12+
13+
@requires_subprocess()
14+
class TestSyncCoordinatorScriptExecution(unittest.TestCase):
15+
"""Tests for how the coordinator executes a target script."""
16+
17+
def run_coordinator(self, cwd, target, *target_args):
18+
"""Execute *target* from *cwd* with the sync coordinator.
19+
20+
Returns the ``(stdout, stderr)`` captured from the coordinator.
21+
"""
22+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
23+
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
24+
server.bind(("127.0.0.1", 0))
25+
server.listen(1)
26+
server.settimeout(SHORT_TIMEOUT)
27+
port = server.getsockname()[1]
28+
29+
cmd = (
30+
sys.executable,
31+
"-m",
32+
"profiling.sampling._sync_coordinator",
33+
str(port),
34+
cwd,
35+
target,
36+
) + tuple(target_args)
37+
process = subprocess.Popen(
38+
cmd,
39+
stdout=subprocess.PIPE,
40+
stderr=subprocess.PIPE,
41+
text=True,
42+
)
43+
try:
44+
conn, _ = server.accept()
45+
with conn:
46+
# _signal_readiness() sends b"ready" before running target.
47+
self.assertEqual(conn.recv(64), b"ready")
48+
stdout, stderr = process.communicate(timeout=SHORT_TIMEOUT)
49+
finally:
50+
if process.poll() is None:
51+
process.kill()
52+
process.wait()
53+
return stdout, stderr
54+
55+
def test_script_directory_on_sys_path(self):
56+
# gh-158540: The coordinator must put the directory containing the
57+
# script (not the working directory) on sys.path, matching
58+
# ``python script.py``, so that modules next to the script can be
59+
# imported even when the coordinator runs from another directory.
60+
with tempfile.TemporaryDirectory() as tmpdir:
61+
script_dir = os.path.join(tmpdir, "sub")
62+
os.mkdir(script_dir)
63+
with open(os.path.join(script_dir, "helper.py"), "w") as f:
64+
f.write("message = 'helper imported'\n")
65+
with open(os.path.join(script_dir, "where.py"), "w") as f:
66+
f.write(
67+
"import os\n"
68+
"import sys\n"
69+
"print('PATH0:', os.path.realpath(sys.path[0]))\n"
70+
"import helper\n"
71+
"print('HELPER:', helper.message)\n"
72+
)
73+
74+
stdout, stderr = self.run_coordinator(
75+
tmpdir, os.path.join("sub", "where.py")
76+
)
77+
78+
self.assertNotIn("ModuleNotFoundError", stderr)
79+
self.assertIn("HELPER: helper imported", stdout)
80+
self.assertIn(f"PATH0: {os.path.realpath(script_dir)}", stdout)
81+
82+
83+
if __name__ == "__main__":
84+
unittest.main()
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Fix :mod:`profiling.sampling` not adding the directory containing the
2+
profiled script to :data:`sys.path`, which prevented the script from
3+
importing modules placed next to it. The script's directory is now added,
4+
matching the behavior of ``python script.py``.

0 commit comments

Comments
 (0)