"""Base classes and registries for data loaders and processors."""
from __future__ import annotations
import importlib
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any
from .bundle import DataBundle
_LOADER_REGISTRY: dict[str, type["BaseDataLoader"]] = {}
_PROCESSOR_REGISTRY: dict[str, type["BaseDataProcessor"]] = {}
# ---------------------------------------------------------------------------
# Loader registry
# ---------------------------------------------------------------------------
[docs]
def register_loader(name: str):
"""Decorator to register a data loader class by name."""
def decorator(cls: type["BaseDataLoader"]) -> type["BaseDataLoader"]:
_LOADER_REGISTRY[name] = cls
cls.name = name
return cls
return decorator
def _ensure_loaders_registered() -> None:
from . import loaders # noqa: F401
def _ensure_processors_registered() -> None:
from . import processors # noqa: F401
[docs]
def get_loader(name: str) -> type["BaseDataLoader"]:
"""Look up a loader by registry name or dotted import path."""
_ensure_loaders_registered()
if name in _LOADER_REGISTRY:
return _LOADER_REGISTRY[name]
if "." in name:
module_path, cls_name = name.rsplit(".", 1)
mod = importlib.import_module(module_path)
cls = getattr(mod, cls_name)
if isinstance(cls, type) and issubclass(cls, BaseDataLoader):
return cls
raise TypeError(f"{name} resolved to {cls!r}, not a BaseDataLoader subclass.")
raise KeyError(
f"Unknown loader {name!r}. "
f"Registered: {sorted(_LOADER_REGISTRY)}. "
f"For external loaders, use a dotted import path."
)
[docs]
def list_loaders() -> list[str]:
"""Return the names of all registered loaders."""
return sorted(_LOADER_REGISTRY)
[docs]
class BaseDataLoader(ABC):
"""Interface for loading readout data from any source.
Subclass this and decorate with ``@register_loader("name")`` to make
the loader discoverable by the pipeline.
Example::
@register_loader("my_format")
class MyLoader(BaseDataLoader):
def load(self, path, **kwargs):
...
return DataBundle(traces=traces, labels=labels)
"""
name: str = "base"
supported_formats: tuple[str, ...] = ()
[docs]
@abstractmethod
def load(self, path: str | Path, **kwargs: Any) -> DataBundle:
"""Load traces and labels from *path*.
Args:
path: File or directory to load from.
**kwargs: Loader-specific options.
Returns:
A DataBundle containing the loaded data.
"""
...
[docs]
@classmethod
def from_config(cls, config: dict[str, Any]) -> "BaseDataLoader":
"""Construct a loader instance from a config dict."""
return cls()
# ---------------------------------------------------------------------------
# Processor registry
# ---------------------------------------------------------------------------
[docs]
def register_processor(name: str):
"""Decorator to register a data processor class by name.
Example::
@register_processor("my_step")
class MyStep(BaseDataProcessor):
def process(self, bundle, **kwargs):
...
return bundle
"""
def decorator(cls: type["BaseDataProcessor"]) -> type["BaseDataProcessor"]:
_PROCESSOR_REGISTRY[name] = cls
cls.name = name
return cls
return decorator
[docs]
def get_processor(name: str) -> type["BaseDataProcessor"]:
"""Look up a processor by registry name or dotted import path."""
_ensure_processors_registered()
if name in _PROCESSOR_REGISTRY:
return _PROCESSOR_REGISTRY[name]
if "." in name:
module_path, cls_name = name.rsplit(".", 1)
mod = importlib.import_module(module_path)
cls = getattr(mod, cls_name)
if isinstance(cls, type) and issubclass(cls, BaseDataProcessor):
return cls
raise TypeError(f"{name} resolved to {cls!r}, not a BaseDataProcessor subclass.")
raise KeyError(
f"Unknown processor {name!r}. "
f"Registered: {sorted(_PROCESSOR_REGISTRY)}. "
f"For external processors, use a dotted import path."
)
[docs]
def list_processors() -> list[str]:
"""Return the names of all registered processors."""
return sorted(_PROCESSOR_REGISTRY)
[docs]
class BaseDataProcessor(ABC):
"""Interface for data processing steps (demod, preprocess, split, etc.).
Subclass this and decorate with ``@register_processor("name")``.
"""
name: str = "base"
requires: list[str] = []
[docs]
@abstractmethod
def process(self, bundle: DataBundle, **kwargs: Any) -> DataBundle:
"""Apply a processing step to the bundle.
Args:
bundle: Current data state.
**kwargs: Processor-specific options.
Returns:
The (possibly modified) DataBundle.
"""
...
# Backward-compatible aliases
BaseLoader = BaseDataLoader
BaseProcessor = BaseDataProcessor