"""Registry mapping ``(component type, backend)`` to implementation classes.""" from collections.abc import Iterator from contextlib import contextmanager from threading import RLock from typing import Callable, TypeVar, cast from .base_component import ComponentMixin from ..enumeration import ComponentType, component_type_name T = TypeVar("T", bound=ComponentMixin) class ComponentRegistry: """Two-level registry: ``component_type -> name -> class``. Supports both direct calls — ``R.register(MyClass, "name")`` — and decorator usage — ``@R.register("name")``. """ def __init__(self) -> None: self._registry: dict[str, dict[str, type[ComponentMixin]]] = {} self._owners: dict[tuple[str, str], str] = {} self._lock = RLock() self._frozen = False def _ensure_mutable(self) -> None: """Reject changes after this registry becomes an immutable template.""" if self._frozen: raise RuntimeError("Component registry is frozen") def _do_register(self, cls: type[T], name: str, *, owner: str | None = None) -> type[T]: """Insert ``cls`` under its component type and reject ambiguous providers.""" try: component_type = component_type_name(getattr(cls, "component_type", None)) except (TypeError, ValueError) as exc: raise TypeError(f"{cls.__name__} must have a non-empty string 'component_type' attribute") from exc if not name: raise ValueError("Component name cannot be empty") with self._lock: self._ensure_mutable() group = self._registry.setdefault(component_type, {}) key = (component_type, name) if name in group: existing = group[name] existing_owner = self._owners[key] new_owner = owner or cls.__module__ if existing is cls and existing_owner == new_owner: return cls raise ValueError( f"Backend '{component_type}:{name}' is provided by both " f"'{existing_owner}' and '{new_owner}'", ) group[name] = cls self._owners[key] = owner or cls.__module__ return cls def add(self, name: str, cls: type[T], *, owner: str) -> type[T]: """Register one explicitly owned plugin contribution.""" return self._do_register(cls, name, owner=owner) def register( self, cls_or_name: type[T] | str, name: str | None = None, ) -> Callable[[type[T]], type[T]] | type[T]: """Register a component class directly, or return a decorator that does so.""" # Direct call: register(MyClass) or register(MyClass, "alias"). if isinstance(cls_or_name, type): cls = cast(type[T], cls_or_name) return self._do_register(cls, name if name is not None else cls.__name__) # Decorator call: @R.register("alias") — must receive a string name. if not isinstance(cls_or_name, str): raise TypeError(f"Expected a class or string, got {type(cls_or_name).__name__}") registration_name = cls_or_name def decorator(decorated_cls: type[T]) -> type[T]: return self._do_register(decorated_cls, registration_name) return decorator def get(self, component_type: ComponentType, name: str) -> type[ComponentMixin] | None: """Look up a registered class; return None if not found.""" with self._lock: return self._registry.get(component_type_name(component_type), {}).get(name) def get_all(self, component_type: ComponentType) -> dict[str, type[ComponentMixin]]: """Return a shallow copy of all classes registered under `component_type`.""" with self._lock: return dict(self._registry.get(component_type_name(component_type), {})) def unregister(self, component_type: ComponentType, name: str) -> bool: """Remove an entry; return True if it existed, False otherwise.""" component_type = component_type_name(component_type) with self._lock: self._ensure_mutable() if (group := self._registry.get(component_type)) and name in group: del group[name] self._owners.pop((component_type, name), None) return True return False def clear(self) -> None: """Drop every registered entry.""" with self._lock: self._ensure_mutable() self._registry.clear() self._owners.clear() def freeze(self) -> None: """Make this registry an immutable template for future copies.""" with self._lock: self._frozen = True @property def frozen(self) -> bool: """Whether mutating operations are disabled.""" with self._lock: return self._frozen def copy(self) -> "ComponentRegistry": """Return an independent registry containing the same providers.""" copied = ComponentRegistry() with self._lock: for component_type, group in self._registry.items(): for name, implementation in group.items(): copied.add(name, implementation, owner=self._owners[(component_type, name)]) return copied @contextmanager def preserve(self, *, allow_mutation: bool = False) -> Iterator[None]: """Restore the registry after code that may register through import side effects.""" with self._lock: registry = {component_type: dict(group) for component_type, group in self._registry.items()} owners = dict(self._owners) frozen = self._frozen if allow_mutation: self._frozen = False try: yield finally: self._registry = registry self._owners = owners self._frozen = frozen # Import-time registry for built-in implementations. Runtime code should use # ``create_application_registry`` rather than mutate this template. R = ComponentRegistry() def create_application_registry() -> ComponentRegistry: """Return a mutable registry initialized from the frozen built-in template.""" return R.copy()