registry.py

"""Decorator-based registry of learning systems.

Architectures add themselves by decorating their class with `@register`.
`discover()` imports every sibling module in `architectures/` so each class
has a chance to register before the CLI looks anything up.
"""

from __future__ import annotations
import importlib
import pkgutil
from typing import Type

from architectures.base import LearningSystem

_REGISTRY: dict[str, Type[LearningSystem]] = {}


def register(cls: Type[LearningSystem]) -> Type[LearningSystem]:
    name = cls.name
    if name in _REGISTRY and _REGISTRY[name] is not cls:
        raise ValueError(f"duplicate architecture name: {name}")
    _REGISTRY[name] = cls
    return cls


def get(name: str) -> Type[LearningSystem]:
    if name not in _REGISTRY:
        choices = ", ".join(sorted(_REGISTRY)) or "(none)"
        raise KeyError(f"unknown architecture '{name}'. available: {choices}")
    return _REGISTRY[name]


def available() -> list[str]:
    return sorted(_REGISTRY)


def discover() -> None:
    """Import every non-base module in `architectures/` so decorators fire."""
    import architectures
    for _, mod_name, _ in pkgutil.iter_modules(architectures.__path__):
        if mod_name.startswith("_") or mod_name == "base":
            continue
        importlib.import_module(f"architectures.{mod_name}")