Skip to content

Commit 0ee04e7

Browse files
committed
Create and use context manager for sys.path augmentation
1 parent 7f388fc commit 0ee04e7

File tree

2 files changed

+26
-6
lines changed

2 files changed

+26
-6
lines changed

astroid/util.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
from __future__ import annotations
77

8+
import contextlib
9+
import sys
810
import warnings
911
from typing import TYPE_CHECKING, Any, Final, Literal
1012

@@ -157,3 +159,24 @@ def safe_infer(
157159
return None # there is some kind of ambiguity
158160
except StopIteration:
159161
return value
162+
163+
def _augment_sys_path(additional_paths: Sequence[str]) -> list[str]:
164+
original = list(sys.path)
165+
changes = []
166+
seen = set()
167+
for additional_path in additional_paths:
168+
if additional_path not in seen:
169+
changes.append(additional_path)
170+
seen.add(additional_path)
171+
172+
sys.path[:] = changes + sys.path
173+
return original
174+
175+
@contextlib.contextmanager
176+
def augmented_sys_path(additional_paths: Sequence[str]) -> Iterator[None]:
177+
"""Augment 'sys.path' by adding entries from additional_paths."""
178+
original = _augment_sys_path(additional_paths)
179+
try:
180+
yield
181+
finally:
182+
sys.path[:] = original

tests/test_modutils.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from astroid import modutils
2323
from astroid.const import PY310_PLUS
2424
from astroid.interpreter._import import spec
25+
from astroid.util import augmented_sys_path
2526

2627
from . import resources
2728

@@ -178,21 +179,19 @@ def test_import_symlink_with_source_outside_of_path(self) -> None:
178179
def test_modpath_from_file_path_order(self) -> None:
179180
"""Test for ordering of paths.
180181
The test does the following:
181-
1. Add a tmp directory to beginning of sys.path
182+
1. Add a tmp directory to beginning of sys.path via augmented_sys_path
182183
2. Create a module file in sub directory of tmp directory
183184
3. If the sub directory is passed as additional directory, module name
184185
should be relative to the subdirectory since additional directory has
185186
higher precedence."""
186-
orig_path = sys.path.copy()
187187
with tempfile.TemporaryDirectory() as tmp_dir:
188-
try:
188+
with augmented_sys_path([tmp_dir]):
189189
mod_name = "module"
190190
sub_dirname = "subdir"
191191
sub_dir = tmp_dir + "/" + sub_dirname
192192
os.mkdir(sub_dir)
193193
module_file = f"{sub_dir}/{mod_name}.py"
194194

195-
sys.path.insert(0, str(tmp_dir))
196195
with open(module_file, "w+", encoding="utf-8"):
197196
pass
198197

@@ -207,8 +206,6 @@ def test_modpath_from_file_path_order(self) -> None:
207206
modutils.modpath_from_file(f"{sub_dir}/{mod_name}.py", [sub_dir]),
208207
[mod_name],
209208
)
210-
finally:
211-
sys.path[:] = orig_path
212209

213210
def test_import_symlink_both_outside_of_path(self) -> None:
214211
with tempfile.NamedTemporaryFile() as tmpfile:

0 commit comments

Comments
 (0)