Yuxuan Yan
81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
"""Registry and factory for daily derived-data plugins."""
|
|
|
|
import importlib
|
|
import importlib.util
|
|
import inspect
|
|
from pathlib import Path
|
|
from typing import Optional, Type
|
|
|
|
from pipeline.derived.base import BaseDerivedData
|
|
|
|
_REGISTRY: dict[str, Type[BaseDerivedData]] = {}
|
|
_builtins_loaded = False
|
|
|
|
|
|
def register_derived(cls: Type[BaseDerivedData]) -> Type[BaseDerivedData]:
|
|
"""Class decorator that registers derived data under ``BaseDerivedData.name``."""
|
|
if not (isinstance(cls, type) and issubclass(cls, BaseDerivedData)):
|
|
raise TypeError(f"{cls!r} is not a BaseDerivedData subclass")
|
|
key = getattr(cls, "name", "")
|
|
if not key:
|
|
raise ValueError(f"{cls.__name__} must set a non-empty class attribute `name`")
|
|
existing = _REGISTRY.get(key)
|
|
if existing is not None and existing is not cls:
|
|
raise ValueError(
|
|
f"Derived data name '{key}' already registered by {existing.__name__}"
|
|
)
|
|
_REGISTRY[key] = cls
|
|
return cls
|
|
|
|
|
|
def available_derived() -> list[str]:
|
|
"""Sorted names of all registered derived-data plugins."""
|
|
_ensure_builtins()
|
|
return sorted(_REGISTRY)
|
|
|
|
|
|
def get_derived(name: str, **params) -> BaseDerivedData:
|
|
"""Instantiate a registered derived-data plugin by name.
|
|
|
|
Only parameters accepted by the plugin class's ``__init__`` are forwarded.
|
|
"""
|
|
_ensure_builtins()
|
|
if name not in _REGISTRY:
|
|
raise KeyError(f"Unknown derived data '{name}'. Available: {sorted(_REGISTRY)}")
|
|
cls = _REGISTRY[name]
|
|
accepted = _accepted_params(cls)
|
|
kwargs = params if accepted is None else {k: v for k, v in params.items() if k in accepted}
|
|
return cls(**kwargs)
|
|
|
|
|
|
def load_derived_module(spec: str) -> None:
|
|
"""Import an external module so its ``@register_derived`` classes register."""
|
|
looks_like_file = spec.endswith(".py") or Path(spec).expanduser().exists()
|
|
if looks_like_file:
|
|
path = Path(spec).expanduser().resolve()
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Derived data module not found: {path}")
|
|
module_spec = importlib.util.spec_from_file_location(path.stem, path)
|
|
if module_spec is None or module_spec.loader is None:
|
|
raise ImportError(f"Cannot load derived data module from {path}")
|
|
module = importlib.util.module_from_spec(module_spec)
|
|
module_spec.loader.exec_module(module)
|
|
else:
|
|
importlib.import_module(spec)
|
|
|
|
|
|
def _accepted_params(cls: Type[BaseDerivedData]) -> Optional[set[str]]:
|
|
"""Param names ``cls.__init__`` accepts, or None if it takes ``**kwargs``."""
|
|
sig = inspect.signature(cls.__init__)
|
|
if any(p.kind is p.VAR_KEYWORD for p in sig.parameters.values()):
|
|
return None
|
|
return {name for name in sig.parameters if name != "self"}
|
|
|
|
|
|
def _ensure_builtins() -> None:
|
|
global _builtins_loaded
|
|
if not _builtins_loaded:
|
|
import pipeline.derived.library # noqa: F401
|
|
_builtins_loaded = True
|
|
|