feat: allow project.chdir() to work in a context (#2516)

src/ape/managers/project.py

@@ -1,5 +1,4 @@
 import json
-import os
 import random
 import shutil
 from collections.abc import Callable, Iterable, Iterator
@@ -32,6 +31,7 @@
 )
 from ape.utils.misc import SOURCE_EXCLUDE_PATTERNS, log_instead_of_fail
 from ape.utils.os import (
+    ChangeDirectory,
     clean_path,
     create_tempdir,
     get_all_files_in_directory,
@@ -2592,22 +2592,50 @@ def chdir(self, path: Optional[Path] = None):
             path (Path): The path of the new project.
               If not given, defaults to the project's path.
         """
-        path = path or self.path
-        os.chdir(path)
-        if self.path == path:
-            return  # Already setup.
+        return ChangeDirectory(
+            self.path,
+            path or self.path,
+            on_push=self._handle_path_changed,
+            on_pop=self._handle_path_restored,
+        )
+
+    def _handle_path_changed(self, path: Path) -> dict:
+        cache: dict = {
+            "__dict__": {},
+            "_session_source_change_check": self._session_source_change_check,
+            "_config_override": self._config_override,
+            "_base_path": self._base_path,
+            "manifest_path": self.manifest_path,
+            "_manifest": self._manifest,
+        }
 
         # New path: clear cached properties.
         for attr in list(self.__dict__.keys()):
             if isinstance(getattr(type(self), attr, None), cached_property):
-                del self.__dict__[attr]
+                cache["__dict__"][attr] = self.__dict__.pop(attr)
 
-        # Re-initialize
         self._session_source_change_check = set()
         self._config_override = {}
         self._base_path = Path(path).resolve()
-        self.manifest_path = self._base_path / ".build" / "__local__.json"
+
+        if self.manifest_path.name == "__local__.json":
+            self.manifest_path = self._base_path / ".build" / "__local__.json"
+
         self._manifest = self.load_manifest()
+        return cache
+
+    def _handle_path_restored(self, cache: dict) -> None:
+        self.__dict__ = {**(self.__dict__ or {}), **cache.get("__dict__", {})}
+        self._session_source_change_check = cache.get("_session_source_change_check", set())
+        self._config_override = cache.get("_config_override", {})
+        if base_path := self._base_path:
+            self._base_path = base_path
+
+        if manifest_path := cache.get("manifest_path"):
+            self.manifest_path = manifest_path
+
+        if manifest := cache.get("_manifest"):
+            self._manifest = manifest
 
     @contextmanager
     def within_project_path(self):

src/ape/utils/os.py

@@ -463,3 +463,49 @@ def within_directory(directory: Path):
     finally:
         if Path.cwd() != here:
             os.chdir(here)
+
+
+class ChangeDirectory:
+    """
+    A context-manager for changing a directory. Initializing it
+    will still change the directory, but you can optionally exit
+    out of it to restore back to the original directory. Additionally,
+    provides hooks to run when each of these events occur.
+    """
+
+    def __init__(
+        self,
+        original_path: Path,
+        new_path: Path,
+        chdir: Optional[Callable[[Path], None]] = None,
+        on_push: Optional[Callable[[Path], dict]] = None,
+        on_pop: Optional[Callable[[dict], None]] = None,
+    ):
+        self.original_path = original_path
+        self.new_path = new_path
+        self._on_push = on_push
+        self._on_pop = on_pop
+        self._chdir = chdir or os.chdir
+
+        # Initiate the change now so you can still use this class
+        # on methods that are not intended to be used in a context.
+        if self.original_path != self.new_path:
+            self._chdir(new_path)
+            self._cache: dict = {} if self._on_push is None else self._on_push(new_path)
+            self._did_change = True
+        else:
+            self._cache = {}
+            self._did_change = False
+
+    def __enter__(self):
+        return self.new_path
+
+    def __exit__(self, *args):
+        if not self._did_change:
+            # Don't do anything. Nothing changed.
+            return
+
+        # Handle the return to the original path.
+        self._chdir(self.original_path)
+        if self._on_pop:
+            self._on_pop(self._cache)

tests/functional/test_project.py

@@ -1065,9 +1065,16 @@ def test_chdir(project):
         project.chdir(new_path)
         assert project.path == new_path
 
-    # Undo.
-    project.chdir(original_path)
-    assert project.path == original_path
+        # Undo.
+        project.chdir(original_path)
+        assert project.path == original_path
+
+        # Show you can also use it in a context.
+        with project.chdir(new_path):
+            assert project.path == new_path
+
+        # It should have automatically undone.
+        assert project.path == original_path
 
 
 def test_within_project_path():