From 0ee04e73e7c6a392d2220a8cff5feeb2c53a32f8 Mon Sep 17 00:00:00 2001
From: akamat10 <akhil.kamat@gmail.com>
Date: Mon, 7 Oct 2024 22:34:37 -0400
Subject: [PATCH] Create and use context manager for sys.path augmentation

---
 astroid/util.py        | 23 +++++++++++++++++++++++
 tests/test_modutils.py |  9 +++------
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/astroid/util.py b/astroid/util.py
index 510b81cc13..3a7432bce2 100644
--- a/astroid/util.py
+++ b/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
diff --git a/tests/test_modutils.py b/tests/test_modutils.py
index 96d0751c44..6b815d986c 100644
--- a/tests/test_modutils.py
+++ b/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: