Update Wall-X to 1.1.0 (#104)

wall_x/config/registry.py

@@ -0,0 +1,99 @@
+"""Runtime registries for typed config variants.
+
+Core Wall-X keeps only base config types in ``wall_x.config``. Optional
+datasets, model families, and optimizer variants register their dataclasses
+from their own packages, so trimmed distributions do not leave dead config
+entries behind.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+_DATA_CONFIG_REGISTRY: dict[str, type[Any]] = {}
+_MODEL_CONFIG_REGISTRY: dict[str, type[Any]] = {}
+_OPTIMIZER_CONFIG_REGISTRY: dict[str, type[Any]] = {}
+_SCHEDULER_CONFIG_REGISTRY: dict[str, type[Any]] = {}
+
+
+def _register(
+    registry: dict[str, type[Any]], names: tuple[str, ...], cls: type[T]
+) -> type[T]:
+    for name in names:
+        if not name:
+            raise ValueError("Config registry name must be non-empty")
+        existing = registry.get(name)
+        if existing is not None and existing is not cls:
+            raise ValueError(
+                f"Config name {name!r} is already registered by "
+                f"{existing.__module__}.{existing.__name__}"
+            )
+        registry[name] = cls
+    return cls
+
+
+def register_data_config(*names: str) -> Callable[[type[T]], type[T]]:
+    return lambda cls: _register(_DATA_CONFIG_REGISTRY, names, cls)
+
+
+def register_model_config(*names: str) -> Callable[[type[T]], type[T]]:
+    return lambda cls: _register(_MODEL_CONFIG_REGISTRY, names, cls)
+
+
+def register_optimizer_config(*names: str) -> Callable[[type[T]], type[T]]:
+    return lambda cls: _register(_OPTIMIZER_CONFIG_REGISTRY, names, cls)
+
+
+def register_scheduler_config(*names: str) -> Callable[[type[T]], type[T]]:
+    return lambda cls: _register(_SCHEDULER_CONFIG_REGISTRY, names, cls)
+
+
+def get_data_config(name: str) -> type[Any] | None:
+    return _DATA_CONFIG_REGISTRY.get(name)
+
+
+def get_model_config(name: str) -> type[Any] | None:
+    return _MODEL_CONFIG_REGISTRY.get(name)
+
+
+def get_optimizer_config(name: str) -> type[Any] | None:
+    return _OPTIMIZER_CONFIG_REGISTRY.get(name)
+
+
+def get_scheduler_config(name: str) -> type[Any] | None:
+    return _SCHEDULER_CONFIG_REGISTRY.get(name)
+
+
+def registered_data_configs() -> list[str]:
+    return sorted(_DATA_CONFIG_REGISTRY)
+
+
+def registered_model_configs() -> list[str]:
+    return sorted(_MODEL_CONFIG_REGISTRY)
+
+
+def registered_optimizer_configs() -> list[str]:
+    return sorted(_OPTIMIZER_CONFIG_REGISTRY)
+
+
+def registered_scheduler_configs() -> list[str]:
+    return sorted(_SCHEDULER_CONFIG_REGISTRY)
+
+
+__all__ = [
+    "get_data_config",
+    "get_model_config",
+    "get_optimizer_config",
+    "get_scheduler_config",
+    "register_data_config",
+    "register_model_config",
+    "register_optimizer_config",
+    "register_scheduler_config",
+    "registered_data_configs",
+    "registered_model_configs",
+    "registered_optimizer_configs",
+    "registered_scheduler_configs",
+]