Remove plugin types, register individual hooks instead

Signed-off-by: Facundo Tuesca <facundo.tuesca@trailofbits.com>

Author: facutuesca

src/pip/_internal/models/plugin.py

@@ -2,17 +2,18 @@
 import logging
 from pathlib import Path
 from types import ModuleType
-from typing import Optional
+from typing import List, Optional
 
 logger = logging.getLogger(__name__)
 
-PLUGIN_TYPE_DIST_INSPECTOR = "dist-inspector"
-SUPPORTED_PLUGIN_TYPES = [PLUGIN_TYPE_DIST_INSPECTOR]
+PLUGIN_HOOK_PRE_DOWNLOAD = "pre_download"
+PLUGIN_HOOK_PRE_EXTRACT = "pre_extract"
+SUPPORTED_PLUGIN_HOOKS = [PLUGIN_HOOK_PRE_DOWNLOAD, PLUGIN_HOOK_PRE_EXTRACT]
 
 
 class Plugin(metaclass=abc.ABCMeta):
     @abc.abstractmethod
-    def plugin_type(self) -> str:
+    def provided_hooks(self) -> List[str]:
         raise NotImplementedError
 
     @property
@@ -21,22 +22,40 @@ def name(self) -> str:
         raise NotImplementedError
 
 
-class DistInspectorPlugin(Plugin):
+class LoadedPlugin(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"
-        ):
+        self._pre_download = None
+        self._pre_extract = None
+        if not hasattr(loaded_module, "provided_hooks"):
             raise ValueError(
-                f'Plugin "{name}" of type {PLUGIN_TYPE_DIST_INSPECTOR} is'
-                "missing pre_download and/or pre_extract definitions"
+                f"Ignoring plugin {name} due to missing provided_hooks definition"
             )
+        for hook in loaded_module.provided_hooks():
+            if hook == PLUGIN_HOOK_PRE_DOWNLOAD:
+                if not hasattr(loaded_module, "pre_download"):
+                    raise ValueError(
+                        f'Plugin "{name}" wants to register a pre-download hook but '
+                        "does not define a pre_download method"
+                    )
+                self._pre_download = loaded_module.pre_download
+            elif hook == PLUGIN_HOOK_PRE_EXTRACT:
+                if not hasattr(loaded_module, "pre_extract"):
+                    raise ValueError(
+                        f'Plugin "{name}" wants to register a pre-extract hook but '
+                        "does not define a pre_extract method"
+                    )
+                self._pre_extract = loaded_module.pre_extract
+            else:
+                raise ValueError(
+                    f'Plugin "{name}" wants to register a hook of unknown type:'
+                    '"{hook}"'
+                )
 
         self._name = name
         self._module = loaded_module
 
-    def plugin_type(self) -> str:
-        return self._module.plugin_type()
+    def provided_hooks(self) -> List[str]:
+        return self._module.provided_hooks()
 
     @property
     def name(self) -> str:
@@ -46,26 +65,19 @@ 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)
+        if self._pre_download is not None:
+            self._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)
+        if self._pre_extract is not None:
+            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)
+def plugin_from_module(name: str, loaded_module: ModuleType) -> Optional[LoadedPlugin]:
+    try:
+        return LoadedPlugin(name, loaded_module)
+    except ValueError as e:
+        logger.warning("Ignoring plugin %s due to error: %s", name, e)
     return None

src/pip/_internal/utils/plugins.py

@@ -4,10 +4,10 @@
 from pathlib import Path
 from typing import Iterator, List
 
-from pip._internal.models.plugin import DistInspectorPlugin, Plugin, plugin_from_module
+from pip._internal.models.plugin import LoadedPlugin, plugin_from_module
 
 logger = logging.getLogger(__name__)
-_loaded_plugins: List[Plugin] = []
+_loaded_plugins: List[LoadedPlugin] = []
 
 
 def iter_entry_points(group_name: str) -> List[EntryPoint]:
@@ -64,8 +64,6 @@ def plugin_pre_download_hook(url: str, filename: str, digest: str) -> None:
     """
 
     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)
 
@@ -82,7 +80,5 @@ def plugin_pre_extract_hook(dist: Path) -> None:
     """
 
     for p in _loaded_plugins:
-        if not isinstance(p, DistInspectorPlugin):
-            continue
         with _only_raise_value_error(p.name):
             p.pre_extract(dist)