Implement provider registry (#633)

* Implement provider registry

* Update src/codemodder/codemods/base_codemod.py

Co-authored-by: Dani Alcala <[email protected]>

---------

Co-authored-by: Dani Alcala <[email protected]>

Author: drdavella

src/codemodder/codemodder.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 from typing import DefaultDict, Sequence
 
-from codemodder import __version__, registry
+from codemodder import __version__, providers, registry
 from codemodder.cli import parse_args
 from codemodder.code_directory import match_files
 from codemodder.codemods.api import BaseCodemod
@@ -142,6 +142,7 @@ def run(original_args) -> int:
     start = datetime.datetime.now()
 
     codemod_registry = registry.load_registered_codemods()
+    provider_registry = providers.load_providers()
 
     # A little awkward, but we need the codemod registry in order to validate potential arguments
     argv = parse_args(original_args, codemod_registry)
@@ -174,6 +175,7 @@ def run(original_args) -> int:
             argv.dry_run,
             argv.verbose,
             codemod_registry,
+            provider_registry,
             repo_manager,
             argv.path_include,
             argv.path_exclude,

src/codemodder/codemods/base_codemod.py

@@ -78,6 +78,7 @@ class BaseCodemod(metaclass=ABCMeta):
     detector: BaseDetector | None
     transformer: BaseTransformerPipeline
     default_extensions: list[str] | None
+    provider: str | None
 
     def __init__(
         self,
@@ -86,12 +87,14 @@ def __init__(
         detector: BaseDetector | None = None,
         transformer: BaseTransformerPipeline,
         default_extensions: list[str] | None = None,
+        provider: str | None = None,
     ):
         # Metadata should only be accessed via properties
         self._metadata = metadata
         self.detector = detector
         self.transformer = transformer
         self.default_extensions = default_extensions or [".py"]
+        self.provider = provider
 
     @property
     @abstractmethod
@@ -159,6 +162,15 @@ def _apply(
         files_to_analyze: list[Path],
         rules: list[str],
     ) -> None:
+        if self.provider and (
+            not (provider := context.providers.get_provider(self.provider))
+            or not provider.is_available
+        ):
+            logger.warning(
+                "provider %s is not available, skipping codemod", self.provider
+            )
+            return
+
         results = (
             # It seems like semgrep doesn't like our fully-specified id format
             self.detector.apply(self.name, context, files_to_analyze)

src/codemodder/codemods/test/utils.py

@@ -7,6 +7,7 @@
 
 from codemodder.context import CodemodExecutionContext
 from codemodder.diff import create_diff
+from codemodder.providers import load_providers
 from codemodder.registry import CodemodCollection, CodemodRegistry
 from codemodder.semgrep import run as semgrep_run
 
@@ -61,6 +62,7 @@ def run_and_assert(
             dry_run=False,
             verbose=False,
             registry=mock.MagicMock(),
+            providers=load_providers(),
             repo_manager=mock.MagicMock(),
             path_include=[f.name for f in files_to_check],
             path_exclude=path_exclude,
@@ -184,6 +186,7 @@ def run_and_assert(
             verbose=False,
             tool_result_files_map={self.tool: [str(tmp_results_file_path)]},
             registry=mock.MagicMock(),
+            providers=load_providers(),
             repo_manager=mock.MagicMock(),
             path_include=[f.name for f in files_to_check],
             path_exclude=path_exclude,

src/codemodder/context.py

@@ -19,6 +19,7 @@
 from codemodder.logging import log_list, logger
 from codemodder.project_analysis.file_parsers.package_store import PackageStore
 from codemodder.project_analysis.python_repo_manager import PythonRepoManager
+from codemodder.providers import ProviderRegistry
 from codemodder.registry import CodemodRegistry
 from codemodder.utils.timer import Timer
 
@@ -37,6 +38,7 @@ class CodemodExecutionContext:
     dry_run: bool = False
     verbose: bool = False
     registry: CodemodRegistry
+    providers: ProviderRegistry
     repo_manager: PythonRepoManager
     timer: Timer
     path_include: list[str]
@@ -51,6 +53,7 @@ def __init__(
         dry_run: bool,
         verbose: bool,
         registry: CodemodRegistry,
+        providers: ProviderRegistry,
         repo_manager: PythonRepoManager,
         path_include: list[str],
         path_exclude: list[str],
@@ -66,6 +69,7 @@ def __init__(
         self._unfixed_findings_by_codemod = {}
         self.dependencies = {}
         self.registry = registry
+        self.providers = providers
         self.repo_manager = repo_manager
         self.timer = Timer()
         self.path_include = path_include

src/codemodder/providers.py

@@ -0,0 +1,60 @@
+from abc import ABCMeta, abstractmethod
+from collections import UserDict
+from importlib.metadata import entry_points
+
+from codemodder.logging import logger
+
+
+class BaseProvider(metaclass=ABCMeta):
+    name: str
+    _resource: str | None
+
+    def __init__(self, name):
+        self.name = name
+        self._resource = self.load()
+
+    @abstractmethod
+    def load(self) -> str | None:
+        pass
+
+    @property
+    def is_available(self) -> bool:
+        return self.resource is not None
+
+    @property
+    def resource(self) -> str:
+        if self._resource is None:
+            raise ValueError(f"Resource for provider {self.name} is not available")
+        return self._resource
+
+
+class ProviderRegistry(UserDict):
+    def add_provider(self, name: str, provider: BaseProvider):
+        self[name] = provider
+
+    def get_provider(self, name: str) -> BaseProvider | None:
+        return self.get(name)
+
+
+def load_providers() -> ProviderRegistry:
+    registry = ProviderRegistry()
+    logger.debug("loading registered providers")
+    for entry_point in entry_points().select(group="codemod_providers"):
+        logger.debug(
+            '- loading provider "%s" from "%s"',
+            entry_point.name,
+            entry_point.module,
+        )
+        try:
+            provider = entry_point.load()
+        except Exception:
+            logger.exception(
+                'Failed to load provider "%s" from "%s": %s',
+                entry_point.name,
+                entry_point.module,
+            )
+            continue
+
+        registry.add_provider(entry_point.name, provider(entry_point.name))
+
+    return registry

tests/codemods/test_xml_transformer.py

@@ -14,14 +14,15 @@ def test_transformer_failure(mocker, caplog):
         side_effect=Exception,
     )
     file_context = FileContext(
-        "home",
+        Path("home"),
         Path("test.xml"),
     )
     execution_context = CodemodExecutionContext(
         directory=mocker.MagicMock(),
         dry_run=True,
         verbose=False,
         registry=mocker.MagicMock(),
+        providers=None,
         repo_manager=mocker.MagicMock(),
         path_include=[],
         path_exclude=[],
@@ -53,6 +54,7 @@ def test_transformer(mocker):
         dry_run=True,
         verbose=False,
         registry=mocker.MagicMock(),
+        providers=None,
         repo_manager=mocker.MagicMock(),
         path_include=[],
         path_exclude=[],

tests/test_context.py

@@ -28,7 +28,9 @@ def test_successful_dependency_description(self, mocker):
         repo_manager = PythonRepoManager(mocker.Mock())
         codemod = registry.match_codemods(codemod_include=["url-sandbox"])[0]
 
-        context = Context(mocker.Mock(), True, False, registry, repo_manager, [], [])
+        context = Context(
+            mocker.Mock(), True, False, registry, None, repo_manager, [], []
+        )
         context.add_dependencies(codemod.id, {Security})
 
         pkg_store_name = "pyproject.toml"
@@ -57,7 +59,9 @@ def test_failed_dependency_description(self, mocker):
         repo_manager = PythonRepoManager(mocker.Mock())
         codemod = registry.match_codemods(codemod_include=["url-sandbox"])[0]
 
-        context = Context(mocker.Mock(), True, False, registry, repo_manager, [], [])
+        context = Context(
+            mocker.Mock(), True, False, registry, None, repo_manager, [], []
+        )
         context.add_dependencies(codemod.id, {Security})
 
         mocker.patch(
@@ -89,6 +93,7 @@ def test_setup_llm_client_no_env_vars(self, mocker):
             True,
             False,
             load_registered_codemods(),
+            None,
             PythonRepoManager(mocker.Mock()),
             [],
             [],
@@ -102,6 +107,7 @@ def test_setup_openai_llm_client(self, mocker):
             True,
             False,
             load_registered_codemods(),
+            None,
             PythonRepoManager(mocker.Mock()),
             [],
             [],
@@ -121,6 +127,7 @@ def test_setup_azure_llm_client(self, mocker):
             True,
             False,
             load_registered_codemods(),
+            None,
             PythonRepoManager(mocker.Mock()),
             [],
             [],
@@ -140,6 +147,7 @@ def test_setup_azure_llm_client_missing_one(self, mocker, env_var):
                 True,
                 False,
                 load_registered_codemods(),
+                None,
                 PythonRepoManager(mocker.Mock()),
                 [],
                 [],
@@ -160,6 +168,7 @@ def test_get_api_version_from_env(self, mocker):
             True,
             False,
             load_registered_codemods(),
+            None,
             PythonRepoManager(mocker.Mock()),
             [],
             [],

tests/test_libcst_transformer.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import mock
 from libcst._exceptions import ParserSyntaxError
 
@@ -55,12 +57,13 @@
 
 
 def _apply_and_assert(mocker, transformer):
-    file_context = FileContext("home", FILE_PATH)
+    file_context = FileContext(Path("home"), FILE_PATH)
     execution_context = CodemodExecutionContext(
         directory=mocker.MagicMock(),
         dry_run=True,
         verbose=False,
         registry=mocker.MagicMock(),
+        providers=None,
         repo_manager=mocker.MagicMock(),
         path_include=[],
         path_exclude=[],
@@ -77,7 +80,7 @@ def _apply_and_assert(mocker, transformer):
 
 def _apply_and_assert_with_tool(mocker, transformer, reason, results):
     file_context = FileContext(
-        "home",
+        Path("home"),
         FILE_PATH,
         results=results,
     )
@@ -86,6 +89,7 @@ def _apply_and_assert_with_tool(mocker, transformer, reason, results):
         dry_run=True,
         verbose=False,
         registry=mocker.MagicMock(),
+        providers=None,
         repo_manager=mocker.MagicMock(),
         path_include=[],
         path_exclude=[],