Create and use context manager for sys.path augmentation

Author: akamat10

astroid/util.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+import contextlib
+import sys
 import warnings
 from typing import TYPE_CHECKING, Any, Final, Literal
 
@@ -157,3 +159,24 @@ def safe_infer(
         return None  # there is some kind of ambiguity
     except StopIteration:
         return value
+
+def _augment_sys_path(additional_paths: Sequence[str]) -> list[str]:
+    original = list(sys.path)
+    changes = []
+    seen = set()
+    for additional_path in additional_paths:
+        if additional_path not in seen:
+            changes.append(additional_path)
+            seen.add(additional_path)
+
+    sys.path[:] = changes + sys.path
+    return original
+
+@contextlib.contextmanager
+def augmented_sys_path(additional_paths: Sequence[str]) -> Iterator[None]:
+    """Augment 'sys.path' by adding entries from additional_paths."""
+    original = _augment_sys_path(additional_paths)
+    try:
+        yield
+    finally:
+        sys.path[:] = original

tests/test_modutils.py

@@ -22,6 +22,7 @@
 from astroid import modutils
 from astroid.const import PY310_PLUS
 from astroid.interpreter._import import spec
+from astroid.util import augmented_sys_path
 
 from . import resources
 
@@ -178,21 +179,19 @@ def test_import_symlink_with_source_outside_of_path(self) -> None:
     def test_modpath_from_file_path_order(self) -> None:
         """Test for ordering of paths.
         The test does the following:
-        1. Add a tmp directory to beginning of sys.path
+        1. Add a tmp directory to beginning of sys.path via augmented_sys_path
         2. Create a module file in sub directory of tmp directory
         3. If the sub directory is passed as additional directory, module name
            should be relative to the subdirectory since additional directory has
            higher precedence."""
-        orig_path = sys.path.copy()
         with tempfile.TemporaryDirectory() as tmp_dir:
-            try:
+            with augmented_sys_path([tmp_dir]):
                 mod_name = "module"
                 sub_dirname = "subdir"
                 sub_dir = tmp_dir + "/" + sub_dirname
                 os.mkdir(sub_dir)
                 module_file = f"{sub_dir}/{mod_name}.py"
 
-                sys.path.insert(0, str(tmp_dir))
                 with open(module_file, "w+", encoding="utf-8"):
                     pass
 
@@ -207,8 +206,6 @@ def test_modpath_from_file_path_order(self) -> None:
                     modutils.modpath_from_file(f"{sub_dir}/{mod_name}.py", [sub_dir]),
                     [mod_name],
                 )
-            finally:
-                sys.path[:] = orig_path
 
     def test_import_symlink_both_outside_of_path(self) -> None:
         with tempfile.NamedTemporaryFile() as tmpfile: