VLA/wall_x/data/_registry.py

"""Data backend registry + multi-verb dispatch.

Each backend registers an entire **module** (not just a single function).
Consumers obtain a backend via :func:`backend_for` and then call any
operation the backend supports (``backend.build`` or optional helper
operations published by that backend). Operations a backend does *not*
implement raise
:class:`MissingOperationError` with a list of backends that do.

Three failure shapes, each pointing at the right next step:

- **Unknown name** (typo / never registered): :class:`KeyError`
  ``Unknown backend 'foo_bar'. Known: [...]``
- **Known but failed to import** (optional dependency missing):
  :class:`RuntimeError`
  ``Backend 'example' failed to import: <ImportError>. Install its
  dependency or switch to one of [...]``
- **Backend exists but operation missing**: :class:`MissingOperationError`
  ``Backend 'example' does not implement 'make_processor'.
  Supported by: ['other_backend']``
"""

from __future__ import annotations

import logging
from types import ModuleType
from typing import Any, Dict

from wall_x.data._bundle import DataBundle
from wall_x.data._protocol import BuildContext

logger = logging.getLogger(__name__)

_BACKENDS: Dict[str, ModuleType] = {}
_import_errors: Dict[str, BaseException] = {}

# The single, process-wide active backend. Set exactly once at config load
# time by :func:`_set_data_backend`; all consumer code reads it via
# :func:`data_backend`. Keeping this as module-level state (rather than
# threading cfg through every callsite) is the whole point of the design.
_DATA_BACKEND: str | None = None


# --- Registration --------------------------------------------------------


def register_module(dataset_type: str, module: ModuleType) -> None:
    """Register a backend module under ``dataset_type``.

    The module must expose a ``build(cfg, ctx) -> DataBundle`` callable.
    Other optional operations are
    discovered lazily via ``getattr`` when ``backend_for(name).<op>`` is
    accessed; backends only declare what they support.
    """
    if not hasattr(module, "build"):
        raise TypeError(
            f"Backend module for {dataset_type!r} must expose a "
            f"``build(cfg, ctx) -> DataBundle`` callable; got {module!r}."
        )
    if dataset_type in _BACKENDS:
        logger.warning("overwriting existing backend registration for %r", dataset_type)
    _BACKENDS[dataset_type] = module


def register(dataset_type: str, build_callable) -> None:
    """Legacy single-callable registration. Wraps ``build_callable`` in a
    minimal module so the new dispatch path still works.

    New backends should prefer :func:`register_module` so they can publish
    multiple operations.
    """
    shim = ModuleType(f"_legacy_backend_shim_{dataset_type}")
    shim.build = build_callable  # type: ignore[attr-defined]
    register_module(dataset_type, shim)


def record_import_error(dataset_type: str, error: BaseException) -> None:
    """Stash the exception that prevented a backend from registering."""
    _import_errors[dataset_type] = error


def available_backends() -> list[str]:
    """Return the names of currently-registered backends."""
    return sorted(_BACKENDS)


# --- Lookup --------------------------------------------------------------


class MissingOperationError(NotImplementedError):
    """Raised when a backend module does not implement a requested op."""

    def __init__(self, backend_name: str, op_name: str) -> None:
        impls = [n for n, mod in _BACKENDS.items() if hasattr(mod, op_name)]
        if impls:
            hint = f"Supported by: {impls}."
        else:
            hint = (
                f"No registered backend implements {op_name!r} - check the "
                f"spelling or add it to the backend module."
            )
        super().__init__(
            f"Backend {backend_name!r} does not implement {op_name!r}. {hint}"
        )
        self.backend_name = backend_name
        self.op_name = op_name


class _BackendProxy:
    """Thin wrapper that forwards ``proxy.<op>(...)`` to the backend module.

    Wrapping (instead of returning the module directly) lets us emit
    ``MissingOperationError`` with a useful "supported by" hint instead
    of plain ``AttributeError``.
    """

    __slots__ = ("_name", "_module")

    def __init__(self, name: str, module: ModuleType) -> None:
        self._name = name
        self._module = module

    def __repr__(self) -> str:
        return f"<BackendProxy name={self._name!r}>"

    def __getattr__(self, op: str):
        attr = getattr(self._module, op, None)
        if attr is None:
            raise MissingOperationError(self._name, op)
        return attr

    def supports(self, op: str) -> bool:
        """Cheap predicate: does this backend implement ``op``?"""
        return hasattr(self._module, op)


def backend_for(name: str) -> _BackendProxy:
    """Look up a backend by ``dataset_type`` name.

    Raises :class:`RuntimeError` (with chained ImportError) if the
    backend is known but failed to import; :class:`KeyError` if the
    name was never registered.
    """
    if name in _BACKENDS:
        return _BackendProxy(name, _BACKENDS[name])
    if name in _import_errors:
        err = _import_errors[name]
        raise RuntimeError(
            f"Backend {name!r} failed to import: {err}. "
            f"Install its dependency or switch to one of "
            f"{available_backends()}."
        ) from err
    raise KeyError(f"Unknown backend {name!r}. Known: {available_backends()}.")


# --- Convenience helpers -------------------------------------------------


def resolve_dataset_type(cfg_or_yaml: Any, default: str = "lerobot") -> str:
    """Pull ``dataset_type`` out of either a typed TrainConfig or a raw yaml dict.

    Lookup order:
    1. ``cfg.data.dataset_type`` (typed TrainConfig)
    2. ``yaml_dict["dataset_type"]`` (legacy flat yaml)
    3. ``default``
    """
    data = getattr(cfg_or_yaml, "data", None)
    if data is not None:
        v = getattr(data, "dataset_type", None)
        if v:
            return v
    if isinstance(cfg_or_yaml, dict):
        v = cfg_or_yaml.get("dataset_type")
        if v:
            return v
    return default


def build_data(cfg: Any, ctx: BuildContext) -> DataBundle:
    """Dispatch ``backend.build(cfg, ctx)`` to the backend named by cfg.

    ``cfg`` is expected to be a typed :class:`TrainConfig` - that's what
    ``wall_x.config.load_config`` returns and what every trainer entry
    point passes in. Raw yaml dicts are not supported here on purpose:
    the typed schema is the contract that gives backends their
    ``cfg.data.dataset_type`` access. If you have a raw dict, use
    :func:`resolve_dataset_type` + :func:`backend_for` directly.
    """
    if not hasattr(cfg, "data") or not hasattr(cfg.data, "dataset_type"):
        raise TypeError(
            f"build_data() expects a typed TrainConfig, got {type(cfg).__name__}. "
            f"Load via wall_x.config.load_config() instead of passing a raw dict."
        )
    return backend_for(cfg.data.dataset_type).build(cfg, ctx)


# --- Active backend (process-global) ------------------------------------


def _set_data_backend(name: str) -> None:
    """Internal - only config loaders should call this.

    Strict semantics: first call wins; same-value re-set is a no-op;
    different-value re-set raises. The error is loud on purpose - it means
    two config loaders disagreed about which backend to use, which is
    almost always a bug (e.g. business code calling this directly, or two
    yamls being loaded in the same process).
    """
    global _DATA_BACKEND
    if _DATA_BACKEND is None:
        _DATA_BACKEND = name
        return
    if _DATA_BACKEND == name:
        return  # idempotent on same value
    raise RuntimeError(
        f"Active backend already set to {_DATA_BACKEND!r}, refusing to "
        f"overwrite with {name!r}. This usually means a config loader was "
        f"called twice with different dataset_type, or business code called "
        f"_set_data_backend directly. Use _reset_data_backend() in tests "
        f"if you need to switch."
    )


def _reset_data_backend() -> None:
    """Internal - clear the active backend. Tests only."""
    global _DATA_BACKEND
    _DATA_BACKEND = None


def has_data_backend() -> bool:
    """Whether a backend has been registered for this process."""
    return _DATA_BACKEND is not None


def data_backend() -> _BackendProxy:
    """Return the active backend proxy. Raises if no config has been loaded.

    Direct construction of ``TrainConfig(...)`` does **not** register a
    backend - that's intentional. Business code must go through
    :func:`wall_x.config.load_config`, which registers the backend before
    returning.
    """
    if _DATA_BACKEND is None:
        raise RuntimeError(
            "No active data backend. Load a TrainConfig via "
            "wall_x.config.load_config() before accessing backend verbs. "
            "Direct TrainConfig() construction does not register a backend."
        )
    return backend_for(_DATA_BACKEND)


__all__ = [
    "register",
    "register_module",
    "record_import_error",
    "available_backends",
    "backend_for",
    "build_data",
    "resolve_dataset_type",
    "MissingOperationError",
    "data_backend",
    "has_data_backend",
]