feat: ✨ PythonModuleLoader

Author: remimd

conftest.py

@@ -1,14 +1,23 @@
 import logging
+from collections.abc import Iterator
+from unittest.mock import patch
 
 import pytest
 
 from injection import Module, mod
 from injection._core.module import Module as CoreModule
+from injection.utils import PythonModuleLoader
 from tests.helpers import EventHistory
 
 logging.basicConfig(level=logging.DEBUG)
 
 
+@pytest.fixture(scope="session", autouse=True)
+def __patch_sys_modules() -> Iterator[None]:
+    with patch.object(PythonModuleLoader, "_sys_modules", {}):
+        yield
+
+
 @pytest.fixture(scope="function", autouse=True)
 def unlock():
     yield

documentation/utils.md

@@ -1,8 +1,8 @@
 # Utils
 
-## load_packages
+## PythonModuleLoader
 
-Useful for put in memory injectables hidden deep within a package. Example:
+Useful for put in memory injectables hidden deep within a package.
 
 ```
 package
@@ -17,6 +17,52 @@ package
 
 To load Injectable1 and Injectable2 into memory you can do the following:
 
+```python
+# Imports
+from injection.utils import PythonModuleLoader
+import package
+```
+
+```python
+def predicate(module_name: str) -> bool:
+    # logic to determine whether the module should be imported or not
+    return True
+
+PythonModuleLoader(predicate).load(package)
+```
+
+### Factory methods
+
+* `from_keywords`
+
+Automatically imports modules whose Python script contains one of the keywords passed in parameter.
+
+```python
+PythonModuleLoader.from_keywords("# Auto-import").load(package)
+```
+
+* `startswith`
+
+Automatically imports modules whose Python script name begins with one of the prefixes passed in parameter.
+
+```python
+profile: str = ...
+PythonModuleLoader.startswith(f"{profile}_").load(package)
+```
+
+* `endswith`
+
+Automatically imports modules whose Python script name ends with one of the suffixes passed in parameter.
+
+```python
+profile: str = ...
+PythonModuleLoader.endswith(f"_{profile}").load(package)
+```
+
+## load_packages
+
+`load_packages` is a simplified version of `PythonModuleLoader`.
+
 ```python
 from injection.utils import load_packages
 

injection/utils.py

@@ -1,15 +1,18 @@
+import itertools
 import sys
-from collections.abc import Callable, Collection, Iterator
+from collections.abc import Callable, Iterator, Mapping
+from dataclasses import dataclass, field
 from importlib import import_module
 from importlib.util import find_spec
+from os.path import isfile
 from pkgutil import walk_packages
+from types import MappingProxyType
 from types import ModuleType as PythonModule
-from typing import ContextManager
+from typing import ClassVar, ContextManager, Self
 
 from injection import Module, mod
-from injection import __name__ as injection_package_name
 
-__all__ = ("load_modules_with_keywords", "load_packages", "load_profile")
+__all__ = ("PythonModuleLoader", "load_packages", "load_profile")
 
 
 def load_profile(*names: str) -> ContextManager[Module]:
@@ -21,76 +24,107 @@ def load_profile(*names: str) -> ContextManager[Module]:
     return mod().load_profile(*names)
 
 
-def load_modules_with_keywords(
+def load_packages(
     *packages: PythonModule | str,
-    keywords: Collection[str] | None = None,
+    predicate: Callable[[str], bool] = lambda module_name: True,
 ) -> dict[str, PythonModule]:
     """
-    Function to import modules from a Python package if one of the keywords is contained in the Python script.
-    The default keywords are:
-        - `from injection `
-        - `from injection.`
-        - `import injection`
+    Function for importing all modules in a Python package.
+    Pass the `predicate` parameter if you want to filter the modules to be imported.
     """
 
-    if keywords is None:
-        keywords = (
-            f"from {injection_package_name} ",
-            f"from {injection_package_name}.",
-            f"import {injection_package_name}",
+    return PythonModuleLoader(predicate).load(*packages).modules
+
+
+@dataclass(repr=False, eq=False, frozen=True, slots=True)
+class PythonModuleLoader:
+    predicate: Callable[[str], bool]
+    __modules: dict[str, PythonModule | None] = field(
+        default_factory=dict,
+        init=False,
+    )
+
+    # To easily mock `sys.modules` in tests
+    _sys_modules: ClassVar[Mapping[str, PythonModule]] = MappingProxyType(sys.modules)
+
+    @property
+    def modules(self) -> dict[str, PythonModule]:
+        return {
+            name: module
+            for name, module in self.__modules.items()
+            if module is not None
+        }
+
+    def load(self, *packages: PythonModule | str) -> Self:
+        modules = itertools.chain.from_iterable(
+            self.__iter_modules(package) for package in packages
         )
+        self.__modules.update(modules)
+        return self
 
-    def predicate(module_name: str) -> bool:
-        spec = find_spec(module_name)
+    def __is_already_loaded(self, module_name: str) -> bool:
+        return any(
+            module_name in modules for modules in (self.__modules, self._sys_modules)
+        )
 
-        if spec and (module_path := spec.origin):
-            with open(module_path, "r") as file:
-                python_script = file.read()
+    def __iter_modules(
+        self,
+        package: PythonModule | str,
+    ) -> Iterator[tuple[str, PythonModule | None]]:
+        if isinstance(package, str):
+            package = import_module(package)
 
-            return bool(python_script) and any(
-                keyword in python_script for keyword in keywords
-            )
+        package_name = package.__name__
 
-        return False
+        try:
+            package_path = package.__path__
+        except AttributeError as exc:
+            raise TypeError(f"`{package_name}` isn't Python package.") from exc
 
-    return load_packages(*packages, predicate=predicate)
+        for info in walk_packages(path=package_path, prefix=f"{package_name}."):
+            name = info.name
 
+            if info.ispkg or self.__is_already_loaded(name):
+                continue
 
-def load_packages(
-    *packages: PythonModule | str,
-    predicate: Callable[[str], bool] = lambda module_name: True,
-) -> dict[str, PythonModule]:
-    """
-    Function for importing all modules in a Python package.
-    Pass the `predicate` parameter if you want to filter the modules to be imported.
-    """
+            module = import_module(name) if self.predicate(name) else None
+            yield name, module
 
-    loaded: dict[str, PythonModule] = {}
+    @classmethod
+    def from_keywords(cls, *keywords: str) -> Self:
+        """
+        Create loader to import modules from a Python package if one of the keywords is
+        contained in the Python script.
+        """
 
-    for package in packages:
-        if isinstance(package, str):
-            package = import_module(package)
+        def predicate(module_name: str) -> bool:
+            spec = find_spec(module_name)
+
+            if spec is None:
+                return False
 
-        loaded |= __iter_modules_from(package, predicate)
+            module_path = spec.origin
 
-    return loaded
+            if module_path is None or not isfile(module_path):
+                return False
 
+            with open(module_path, "r") as script:
+                return any(keyword in line for line in script for keyword in keywords)
 
-def __iter_modules_from(
-    package: PythonModule,
-    predicate: Callable[[str], bool],
-) -> Iterator[tuple[str, PythonModule]]:
-    package_name = package.__name__
+        return cls(predicate)
 
-    try:
-        package_path = package.__path__
-    except AttributeError as exc:
-        raise TypeError(f"`{package_name}` isn't Python package.") from exc
+    @classmethod
+    def startswith(cls, *prefixes: str) -> Self:
+        def predicate(module_name: str) -> bool:
+            script_name = module_name.split(".")[-1]
+            return any(script_name.startswith(prefix) for prefix in prefixes)
 
-    for info in walk_packages(path=package_path, prefix=f"{package_name}."):
-        name = info.name
+        return cls(predicate)
 
-        if info.ispkg or name in sys.modules or not predicate(name):
-            continue
+    @classmethod
+    def endswith(cls, *suffixes: str) -> Self:
+        def predicate(module_name: str) -> bool:
+            script_name = module_name.split(".")[-1]
+            return any(script_name.endswith(suffix) for suffix in suffixes)
 
-        yield name, import_module(name)
+        return cls(predicate)