feat: Abstract out caching in discovery (#2946)

Author: esafak

docs/changelog/2074.feature.rst

@@ -0,0 +1 @@
+Abstract out caching in discovery - by :user:`esafak`.

pyproject.toml

@@ -47,6 +47,7 @@ dependencies = [
   "filelock>=3.12.2,<4",
   "importlib-metadata>=6.6; python_version<'3.8'",
   "platformdirs>=3.9.1,<5",
+  "typing-extensions>=4.13.2; python_version<'3.11'",
 ]
 optional-dependencies.docs = [
   "furo>=2023.7.26",

src/virtualenv/discovery/cache.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+try:
+    from typing import Self  # pragma: ≥ 3.11 cover
+except ImportError:
+    from typing_extensions import Self  # pragma: < 3.11 cover
+
+
+class Cache(ABC):
+    """
+    A generic cache interface.
+
+    Add a close() method if the cache needs to perform any cleanup actions,
+    and an __exit__ method to allow it to be used in a context manager.
+    """
+
+    @abstractmethod
+    def get(self, key: str) -> Any | None:
+        """
+        Get a value from the cache.
+
+        :param key: the key to retrieve
+        :return: the cached value, or None if not found
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def set(self, key: str, value: Any) -> None:
+        """
+        Set a value in the cache.
+
+        :param key: the key to set
+        :param value: the value to cache
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def remove(self, key: str) -> None:
+        """
+        Remove a value from the cache.
+
+        :param key: the key to remove
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def clear(self) -> None:
+        """Clear the entire cache."""
+        raise NotImplementedError
+
+    def __enter__(self) -> Self:
+        return self
+
+
+__all__ = [
+    "Cache",
+]

src/virtualenv/discovery/cached_py_info.py

@@ -7,7 +7,6 @@
 
 from __future__ import annotations
 
-import hashlib
 import logging
 import os
 import random
@@ -17,8 +16,14 @@
 from shlex import quote
 from string import ascii_lowercase, ascii_uppercase, digits
 from subprocess import Popen
+from typing import TYPE_CHECKING
 
-from virtualenv.app_data import AppDataDisabled
+from virtualenv.app_data.na import AppDataDisabled
+from virtualenv.discovery.file_cache import FileCache
+
+if TYPE_CHECKING:
+    from virtualenv.app_data.base import AppData
+    from virtualenv.discovery.cache import Cache
 from virtualenv.discovery.py_info import PythonInfo
 from virtualenv.util.subprocess import subprocess
 
@@ -27,9 +32,20 @@
 LOGGER = logging.getLogger(__name__)
 
 
-def from_exe(cls, app_data, exe, env=None, raise_on_error=True, ignore_cache=False):  # noqa: FBT002, PLR0913
+def from_exe(  # noqa: PLR0913
+    cls,
+    app_data,
+    exe,
+    env=None,
+    *,
+    raise_on_error=True,
+    ignore_cache=False,
+    cache: Cache | None = None,
+) -> PythonInfo | None:
     env = os.environ if env is None else env
-    result = _get_from_cache(cls, app_data, exe, env, ignore_cache=ignore_cache)
+    if cache is None:
+        cache = FileCache(app_data)
+    result = _get_from_cache(cls, app_data, exe, env, cache, ignore_cache=ignore_cache)
     if isinstance(result, Exception):
         if raise_on_error:
             raise result
@@ -38,63 +54,35 @@ def from_exe(cls, app_data, exe, env=None, raise_on_error=True, ignore_cache=Fal
     return result
 
 
-def _get_from_cache(cls, app_data, exe, env, ignore_cache=True):  # noqa: FBT002
+def _get_from_cache(cls, app_data: AppData, exe: str, env, cache: Cache, *, ignore_cache: bool) -> PythonInfo:  # noqa: PLR0913
     # note here we cannot resolve symlinks, as the symlink may trigger different prefix information if there's a
     # pyenv.cfg somewhere alongside on python3.5+
     exe_path = Path(exe)
     if not ignore_cache and exe_path in _CACHE:  # check in the in-memory cache
         result = _CACHE[exe_path]
     else:  # otherwise go through the app data cache
-        py_info = _get_via_file_cache(cls, app_data, exe_path, exe, env)
-        result = _CACHE[exe_path] = py_info
+        result = _CACHE[exe_path] = _get_via_file_cache(cls, app_data, exe_path, exe, env, cache)
     # independent if it was from the file or in-memory cache fix the original executable location
     if isinstance(result, PythonInfo):
         result.executable = exe
     return result
 
 
-def _get_via_file_cache(cls, app_data, path, exe, env):
-    path_text = str(path)
-    try:
-        path_modified = path.stat().st_mtime
-    except OSError:
-        path_modified = -1
-    py_info_script = Path(os.path.abspath(__file__)).parent / "py_info.py"
-    try:
-        py_info_hash = hashlib.sha256(py_info_script.read_bytes()).hexdigest()
-    except OSError:
-        py_info_hash = None
-
-    if app_data is None:
-        app_data = AppDataDisabled()
-    py_info, py_info_store = None, app_data.py_info(path)
-    with py_info_store.locked():
-        if py_info_store.exists():  # if exists and matches load
-            data = py_info_store.read()
-            of_path = data.get("path")
-            of_st_mtime = data.get("st_mtime")
-            of_content = data.get("content")
-            of_hash = data.get("hash")
-            if of_path == path_text and of_st_mtime == path_modified and of_hash == py_info_hash:
-                py_info = cls._from_dict(of_content.copy())
-                sys_exe = py_info.system_executable
-                if sys_exe is not None and not os.path.exists(sys_exe):
-                    py_info_store.remove()
-                    py_info = None
-            else:
-                py_info_store.remove()
-        if py_info is None:  # if not loaded run and save
-            failure, py_info = _run_subprocess(cls, exe, app_data, env)
-            if failure is None:
-                data = {
-                    "st_mtime": path_modified,
-                    "path": path_text,
-                    "content": py_info._to_dict(),  # noqa: SLF001
-                    "hash": py_info_hash,
-                }
-                py_info_store.write(data)
-            else:
-                py_info = failure
+def _get_via_file_cache(cls, app_data: AppData, path: Path, exe: str, env, cache: Cache) -> PythonInfo:  # noqa: PLR0913
+    py_info = cache.get(path)
+    if py_info is not None:
+        py_info = cls._from_dict(py_info)
+        sys_exe = py_info.system_executable
+        if sys_exe is not None and not os.path.exists(sys_exe):
+            cache.remove(path)
+            py_info = None
+
+    if py_info is None:  # if not loaded run and save
+        failure, py_info = _run_subprocess(cls, exe, app_data, env)
+        if failure is None:
+            cache.set(path, py_info._to_dict())  # noqa: SLF001
+        else:
+            py_info = failure
     return py_info
 
 
@@ -120,6 +108,8 @@ def _run_subprocess(cls, exe, app_data, env):
 
     start_cookie = gen_cookie()
     end_cookie = gen_cookie()
+    if app_data is None:
+        app_data = AppDataDisabled()
     with app_data.ensure_extracted(py_info_script) as py_info_script:
         cmd = [exe, str(py_info_script), start_cookie, end_cookie]
         # prevent sys.prefix from leaking into the child process - see https://bugs.python.org/issue22490
@@ -182,8 +172,12 @@ def __repr__(self) -> str:
         return cmd_repr
 
 
-def clear(app_data):
-    app_data.py_info_clear()
+def clear(app_data=None, cache=None):
+    """Clear the cache."""
+    if cache is None and app_data is not None:
+        cache = FileCache(app_data)
+    if cache is not None:
+        cache.clear()
     _CACHE.clear()
 
 

src/virtualenv/discovery/file_cache.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import hashlib
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from virtualenv.app_data.na import AppDataDisabled
+from virtualenv.discovery.cache import Cache
+
+if TYPE_CHECKING:
+    from virtualenv.app_data.base import AppData
+
+LOGGER = logging.getLogger(__name__)
+
+
+class FileCache(Cache):
+    def __init__(self, app_data: AppData) -> None:
+        self.app_data = app_data if app_data is not None else AppDataDisabled()
+
+    def get(self, key: Path):
+        """Get a value from the file cache."""
+        py_info, py_info_store = None, self.app_data.py_info(key)
+        with py_info_store.locked():
+            if py_info_store.exists():
+                py_info = self._read_from_store(py_info_store, key)
+        return py_info
+
+    def set(self, key: Path, value: dict) -> None:
+        """Set a value in the file cache."""
+        py_info_store = self.app_data.py_info(key)
+        with py_info_store.locked():
+            path_text = str(key)
+            try:
+                path_modified = key.stat().st_mtime
+            except OSError:
+                path_modified = -1
+
+            py_info_script = Path(__file__).parent / "py_info.py"
+            try:
+                py_info_hash = hashlib.sha256(py_info_script.read_bytes()).hexdigest()
+            except OSError:
+                py_info_hash = None
+
+            data = {
+                "st_mtime": path_modified,
+                "path": path_text,
+                "content": value,
+                "hash": py_info_hash,
+            }
+            py_info_store.write(data)
+
+    def remove(self, key: Path) -> None:
+        """Remove a value from the file cache."""
+        py_info_store = self.app_data.py_info(key)
+        with py_info_store.locked():
+            if py_info_store.exists():
+                py_info_store.remove()
+
+    def clear(self) -> None:
+        """Clear the entire file cache."""
+        self.app_data.py_info_clear()
+
+    def _read_from_store(self, py_info_store, path: Path):
+        data = py_info_store.read()
+        path_text = str(path)
+        try:
+            path_modified = path.stat().st_mtime
+        except OSError:
+            path_modified = -1
+
+        py_info_script = Path(__file__).parent / "py_info.py"
+        try:
+            py_info_hash = hashlib.sha256(py_info_script.read_bytes()).hexdigest()
+        except OSError:
+            py_info_hash = None
+
+        of_path = data.get("path")
+        of_st_mtime = data.get("st_mtime")
+        of_content = data.get("content")
+        of_hash = data.get("hash")
+
+        if of_path == path_text and of_st_mtime == path_modified and of_hash == py_info_hash:
+            return of_content
+
+        py_info_store.remove()
+        return None
+
+
+__all__ = [
+    "FileCache",
+]