feat: allow keeping stuff when isolating data dir (#2515)

Author: antazoey

src/ape/managers/config.py

@@ -1,9 +1,10 @@
 import os
-from collections.abc import Iterator
+import shutil
+from collections.abc import Iterable, Iterator
 from contextlib import contextmanager
 from functools import cached_property
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from ape.api.config import ApeConfig
 from ape.managers.base import BaseManager
@@ -113,20 +114,41 @@ def extract_config(cls, manifest: "PackageManifest", **overrides) -> ApeConfig:
         return ApeConfig.from_manifest(manifest, **overrides)
 
     @contextmanager
-    def isolate_data_folder(self) -> Iterator[Path]:
+    def isolate_data_folder(
+        self, keep: Optional[Union[Iterable[str], str]] = None
+    ) -> Iterator[Path]:
         """
         Change Ape's DATA_FOLDER to point a temporary path,
         in a context, for testing purposes. Any data
         cached to disk will not persist.
+
+        Args:
+            keep (Optional[Union[Iterable[str], str]]): Optionally, pass in
+              a key of subdirectory names to include in the new isolated
+              data folder. For example, pass ing ``"packages"`` to avoid
+              having to re-download dependencies in an isolated environment.
+
+        Returns:
+            Iterator[Path]: The temporary data folder.
         """
         original_data_folder = self.DATA_FOLDER
         if in_tempdir(original_data_folder):
             # Already isolated.
             yield original_data_folder
 
         else:
+            keep = [keep] if isinstance(keep, str) else keep or []
             try:
                 with create_tempdir() as temp_data_folder:
+                    # Copy in items from "keep".
+                    for item in keep:
+                        path_to_keep = original_data_folder / item
+                        if not path_to_keep.is_dir():
+                            continue
+
+                        dest_path = temp_data_folder / item
+                        shutil.copytree(path_to_keep, dest_path)
+
                     self.DATA_FOLDER = temp_data_folder
                     yield temp_data_folder
 

tests/functional/test_config.py

@@ -1,5 +1,6 @@
 import os
 import re
+import shutil
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
@@ -689,3 +690,44 @@ def test_project_level_settings(project):
         assert project.config.my_string == "my_string"
         assert project.config.my_int == 123
         assert project.config.my_bool is True
+
+
+def test_isolate_data_folder(config):
+    original_data_folder = config.DATA_FOLDER
+    madeup_path = Path.home() / ".ape" / "__madeup__"
+    config.DATA_FOLDER = madeup_path
+    try:
+        with config.isolate_data_folder():
+            assert config.DATA_FOLDER != original_data_folder
+    finally:
+        config.DATA_FOLDER = original_data_folder
+
+
+def test_isolate_data_folder_already_isolated(config):
+    data_folder = config.DATA_FOLDER
+    with config.isolate_data_folder():
+        # Already isolated, so there is no change.
+        assert config.DATA_FOLDER == data_folder
+
+
+def test_isolate_data_folder_keep(config):
+    original_data_folder = config.DATA_FOLDER
+    madeup_path = Path.home() / ".ape" / "__core_aoe_test__"
+    madeup_path.mkdir(parents=True, exist_ok=True)
+    sub_dir = madeup_path / "subdir"
+    file = sub_dir / "file.txt"
+    sub_dir.mkdir(parents=True, exist_ok=True)
+    file.write_text("search for 'test_isolate_data_folder_keep' ape's tests.")
+
+    config.DATA_FOLDER = madeup_path
+    try:
+        with config.isolate_data_folder(keep="subdir"):
+            assert config.DATA_FOLDER != original_data_folder
+
+        expected_file = config.DATA_FOLDER / "subdir" / "file.txt"
+        assert expected_file.is_file()
+
+    finally:
+        config.DATA_FOLDER = original_data_folder
+        if madeup_path.is_dir():
+            shutil.rmtree(str(madeup_path), ignore_errors=True)