Add initial support for plugins

pypa · facutuesca · Oct 1, 2024 · Oct 1, 2024 · Oct 1, 2024 · Oct 29, 2024
commit 2cb0ca5b35e881a3703552ea9a897da39a2ac2fb
diff --git a/src/pip/_internal/models/plugin.py b/src/pip/_internal/models/plugin.py
@@ -0,0 +1,71 @@
+import abc
+import logging
+from pathlib import Path
+from types import ModuleType
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+PLUGIN_TYPE_DIST_INSPECTOR = "dist-inspector"
+SUPPORTED_PLUGIN_TYPES = [PLUGIN_TYPE_DIST_INSPECTOR]
+
+
+class Plugin(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def plugin_type(self) -> str:
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def name(self) -> str:
+        raise NotImplementedError
+
+
+class DistInspectorPlugin(Plugin):
+    def __init__(self, name: str, loaded_module: ModuleType):
+        assert loaded_module.plugin_type() == PLUGIN_TYPE_DIST_INSPECTOR
+        if not hasattr(loaded_module, "pre_download") or not hasattr(
+            loaded_module, "pre_extract"
+        ):
+            raise ValueError(
+                f'Plugin "{name}" of type {PLUGIN_TYPE_DIST_INSPECTOR} is'
+                "missing pre_download and/or pre_extract definitions"
+            )
+
+        self._name = name
+        self._module = loaded_module
+
+    def plugin_type(self) -> str:
+        return self._module.plugin_type()
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def pre_download(self, url: str, filename: str, digest: str) -> None:
+        # contract: `pre_download` raises `ValueError` to terminate
+        # the operation that intends to download `filename` from `url`
+        # with hash `digest`
+        self._module.pre_download(url=url, filename=filename, digest=digest)
+
+    def pre_extract(self, dist: Path) -> None:
+        # contract: `pre_extract` raises `ValueError` to terminate
+        # the operation that intends to unarchive `dist`
+        self._module.pre_extract(dist)
+
+
+def plugin_from_module(name: str, loaded_module: ModuleType) -> Optional[Plugin]:
+    if not hasattr(loaded_module, "plugin_type"):
+        logger.warning("Ignoring plugin %s due to missing plugin_type definition", name)
+    plugin_type = loaded_module.plugin_type()
+    if plugin_type not in SUPPORTED_PLUGIN_TYPES:
+        logger.warning(
+            "Ignoring plugin %s due to unknown plugin type: %s", name, plugin_type
+        )
+
+    if plugin_type == PLUGIN_TYPE_DIST_INSPECTOR:
+        try:
+            return DistInspectorPlugin(name, loaded_module)
+        except ValueError as e:
+            logger.warning("Ignoring plugin %s due to error: %s", name, e)
+    return None
diff --git a/src/pip/_internal/operations/install/wheel.py b/src/pip/_internal/operations/install/wheel.py
@@ -15,6 +15,7 @@
 from base64 import urlsafe_b64encode
 from email.message import Message
 from itertools import chain, filterfalse, starmap
+from pathlib import Path
 from typing import (
     IO,
        TYPE_CHECKING,
@@ -52,6 +53,7 @@
 from pip._internal.models.scheme import SCHEME_KEYS, Scheme
 from pip._internal.utils.filesystem import adjacent_tmp_file, replace
 from pip._internal.utils.misc import StreamWrapper, ensure_dir, hash_file, partition
+from pip._internal.utils.plugins import plugin_pre_extract_hook
 from pip._internal.utils.unpacking import (
     current_umask,
     is_within_directory,
@@ -727,6 +729,12 @@ def install_wheel(
     direct_url: Optional[DirectUrl] = None,
     requested: bool = False,
 ) -> None:
+    try:
+        plugin_pre_extract_hook(Path(wheel_path))
+    except ValueError as e:
+        raise InstallationError(
+            f"Could not unpack file {wheel_path} due to plugin:\n{e}"
+        )
     with ZipFile(wheel_path, allowZip64=True) as z:
         with req_error_context(req_description):
             _install_wheel(

diff --git a/src/pip/_internal/operations/prepare.py b/src/pip/_internal/operations/prepare.py
@@ -50,6 +50,7 @@
     hide_url,
     redact_auth_from_requirement,
 )
+from pip._internal.utils.plugins import plugin_pre_download_hook
 from pip._internal.utils.temp_dir import TempDirectory
 from pip._internal.utils.unpacking import unpack_file
 from pip._internal.vcs import vcs
@@ -461,6 +462,14 @@ def _complete_partial_requirements(
         for req in partially_downloaded_reqs:
             assert req.link
             links_to_fully_download[req.link] = req
+            try:
+                plugin_pre_download_hook(
+                    url=req.link.url, filename=req.link.filename, digest=req.link.hash
+                )
+            except ValueError as e:
+                raise InstallationError(
+                    f"Could not install requirement {req} due to plugin:\n{e}"
+                )
 
         batch_download = self._batch_download(
             links_to_fully_download.keys(),
@@ -595,6 +604,9 @@ def _prepare_linked_requirement(
             local_file = None
         elif link.url not in self._downloaded:
             try:
+                plugin_pre_download_hook(
+                    url=req.link.url, filename=req.link.filename, digest=req.link.hash
+                )
                 local_file = unpack_url(
                     link,
                     req.source_dir,
@@ -608,6 +620,10 @@ def _prepare_linked_requirement(
                     f"Could not install requirement {req} because of HTTP "
                     f"error {exc} for URL {link}"
                 )
+            except ValueError as e:
+                raise InstallationError(
+                    f"Could not install requirement {req} due to plugin:\n{e}"
+                )
         else:
             file_path = self._downloaded[link.url]
             if hashes:

diff --git a/src/pip/_internal/utils/plugins.py b/src/pip/_internal/utils/plugins.py
@@ -0,0 +1,72 @@
+import contextlib
+import logging
+from pathlib import Path
+from typing import Iterator, List
+
+from pip._vendor.pygments.plugin import iter_entry_points
+
+from pip._internal.models.plugin import DistInspectorPlugin, Plugin, plugin_from_module
+
+logger = logging.getLogger(__name__)
+
+_loaded_plugins: List[Plugin] = []
+for entrypoint in iter_entry_points(group_name="pip.plugins"):
+    try:
+        module = entrypoint.load()
+    except ModuleNotFoundError:
+        logger.warning("Tried to load plugin %s but failed", entrypoint.name)
+        continue
+    plugin = plugin_from_module(entrypoint.name, module)
+    if plugin is not None:
+        _loaded_plugins.append(plugin)
+
+
+@contextlib.contextmanager
+def _only_raise_value_error(plugin_name: str) -> Iterator[None]:
+    try:
+        yield
+    except ValueError as e:
+        raise ValueError(f"Plugin {plugin_name}: {e}") from e
+    except Exception as e:
+        logger.warning(
+            "Plugin %s raised an unexpected exception type: %s",
+            plugin_name,
+            {e.__class__.__name__},
+        )
+        raise ValueError(f"Plugin {plugin_name}: {e}") from e
+
+
+def plugin_pre_download_hook(url: str, filename: str, digest: str) -> None:
+    """Call the pre-download hook of all loaded plugins
+
+    This function should be called right before a distribution is downloaded.
+    It will go through all the loaded plugins and call their `pre_download(url)`
+    function.
+    Only ValueError will be raised. If the plugin (incorrectly) raises another
+    exception type, this function will wrap it as a ValueError and log
+    a warning.
+    """
+
+    for p in _loaded_plugins:
+        if not isinstance(p, DistInspectorPlugin):
+            continue
+        with _only_raise_value_error(p.name):
+            p.pre_download(url=url, filename=filename, digest=digest)
+
+
+def plugin_pre_extract_hook(dist: Path) -> None:
+    """Call the pre-extract hook of all loaded plugins
+
+    This function should be called right before a distribution is extracted.
+    It will go through all the loaded plugins and call their `pre_extract(dist)`
+    function.
+    Only ValueError will be raised. If the plugin (incorrectly) raises another
+    exception type, this function will wrap it as a ValueError and log
+    a warning.
+    """
+
+    for p in _loaded_plugins:
+        if not isinstance(p, DistInspectorPlugin):
+            continue
+        with _only_raise_value_error(p.name):
+            p.pre_extract(dist)
diff --git a/src/pip/_internal/utils/unpacking.py b/src/pip/_internal/utils/unpacking.py
@@ -8,6 +8,7 @@
 import sys
 import tarfile
 import zipfile
+from pathlib import Path
 from typing import Iterable, List, Optional
 from zipfile import ZipInfo
 
@@ -19,6 +20,7 @@
     ZIP_EXTENSIONS,
 )
 from pip._internal.utils.misc import ensure_dir
+from pip._internal.utils.plugins import plugin_pre_extract_hook
 
 logger = logging.getLogger(__name__)
 
@@ -312,6 +314,11 @@ def unpack_file(
     content_type: Optional[str] = None,
 ) -> None:
     filename = os.path.realpath(filename)
+    try:
+        plugin_pre_extract_hook(Path(filename))
+    except ValueError as e:
+        raise InstallationError(f"Could not unpack file {filename} due to plugin:\n{e}")
+
     if (
         content_type == "application/zip"
         or filename.lower().endswith(ZIP_EXTENSIONS)