Better loader discovery

CLAUDE.md

@@ -113,7 +113,7 @@ The code lives under `src/lib/` and is organized around three concepts: **source
 ### Data flow
 
 1. `parsing.get_parsed_args()` (`src/lib/parsing/parse.py`) builds a flat argparse parser with positional `prefix` (choices-constrained to known prefixes) and optional `variable`, plus all registered adaptor/hook flags. Returns an `Args` namespace (`src/lib/parsing/args.py`) directly.
-2. `args.get_animation()` looks up the loader class for `prefix` in `LOADERS` (from `src/lib/data/loader_registry.py`) and constructs it with `(prefix, variable)`. Each loader handles step discovery and data loading internally. It then calls `compile_source` (`src/lib/data/compile.py`) to wrap it in a `DataSourceWithPipeline` made of the user-supplied `Adaptor` list. If no `Versus` adaptor is present, a default one is appended (`y,z` vs `t`) — this is what selects axes and time dim.
+2. `parsing.get_parsed_args()` also calls `discover_loaders(CONFIG.data_dir)` (`src/lib/data/loader.py`) to map each prefix to its loader class, then instantiates the chosen loader as `args.loader`. `args.get_animation()` calls `compile_source` (`src/lib/data/compile.py`) to wrap `args.loader` in a `DataSourceWithPipeline` made of the user-supplied `Adaptor` list. If no `Versus` adaptor is present, a default one is appended (`y,z` vs `t`) — this is what selects axes and time dim.
 3. `source.get_data()` loads raw data and runs the pipeline, returning a `DataWithAttrs` (a `Field` wrapping `xr.Dataset`, or a `List` wrapping `pd.DataFrame` / `dd.DataFrame`). `DataSource` is an ABC with a single abstract method `get_data()`.
 4. `get_plot(data)` (`src/lib/plotting/get_plot.py`) picks a `Renderer` based on data type and geometry (`Field1dRenderer`, `Field2dRenderer`, `PolarFieldRenderer`, `ScatterRenderer`), then wraps it in `StaticPlot` or `AnimatedPlot` based on whether `time_dim` is present. `ScatterRenderer` requires exactly 2 `spatial_dims`. Each renderer defines `make_init_data`/`init` (one-time setup with frame 0) and `make_update_data`/`draw` (per-frame update for animations). Hooks are added to the plot object after construction.
 5. Hooks (`src/lib/plotting/hooks/`) such as `--scale log`, `--grid`, `--vline`, `--fit` are appended onto the chosen plot before `show()`/`save()`.
@@ -122,7 +122,9 @@ The code lives under `src/lib/` and is organized around three concepts: **source
 
 `src/lib/__init__.py` imports `lib.data.loaders`, `lib.data.adaptors`, and `lib.plotting.hooks`, whose `__init__.py` files glob and `importlib.import_module` every sibling `*.py`.
 
-**Loaders** (`src/lib/data/loaders/`) register themselves via the `@loader(...)` decorator from `src/lib/data/loader_registry.py`, which populates `LOADERS: dict[str, type[DataSource]]`. Each loader class takes `(prefix: str, active_key: str | None)` and handles step discovery and data loading internally. `parse._get_parser()` uses `LOADERS.keys()` as the valid prefix choices. To add a new data source, drop a file into `src/lib/data/loaders/` and decorate the class with `@loader("prefix1", "prefix2", ...)`.
+**Loaders** (`src/lib/data/loaders/`) register themselves via the bare `@loader` decorator from `src/lib/data/loader.py`, which appends to `LOADERS: list[type[Loader]]`. Each loader class exposes a `discover_prefixes(cls, data_dir: Path) -> list[str]` classmethod that returns the prefixes it claims in that directory.
+
+`discover_loaders(data_dir)` polls every registered loader's `discover_prefixes()` and returns a `dict[str, type[Loader]]`; the resulting keys become argparse's prefix choices. On prefix conflicts, the later-registered loader wins with a `UserWarning`, so user-defined loaders can shadow built-ins. To add a new data source, drop a file into `src/lib/data/loaders/`, decorate the class with `@loader`, and implement `discover_prefixes`.
 
 **Adaptors and hooks** register their argparse flags via the `@arg_parser(...)` / `@const_arg(...)` decorators in `src/lib/parsing/args_registry.py`, which append to the module-level `CUSTOM_ARGS` list. `parse._get_parser()` then iterates `CUSTOM_ARGS` and adds them to the parser. To add a new adaptor or hook, drop a new file into `src/lib/data/adaptors/` (or `src/lib/plotting/hooks/`) and decorate its parse function.
 

src/lib/data/compile.py

@@ -1,7 +1,7 @@
 from lib.data.adaptor import Adaptor
 from lib.data.adaptors.versus import Versus
+from lib.data.data_source import DataSource, DataSourceWithPipeline
 from lib.data.pipeline import Pipeline
-from lib.data.source import DataSource, DataSourceWithPipeline
 
 
 def compile_source(loader: DataSource, adaptors: list[Adaptor]) -> DataSource:

src/lib/data/loader.py

@@ -0,0 +1,56 @@
+import warnings
+from abc import abstractmethod
+from pathlib import Path
+
+from lib.data.data_source import DataSource
+from lib.file_util import get_available_steps
+
+
+class Loader(DataSource):
+    @classmethod
+    @abstractmethod
+    def discover_prefixes(cls, data_dir: Path) -> list[str]:
+        """Return prefixes this loader can handle in data_dir."""
+
+    @classmethod
+    @abstractmethod
+    def suffix(cls) -> str:
+        """Return the suffix that this loader supports."""
+
+    def __init__(self, prefix: str, active_key: str | None = None):
+        self.prefix = prefix
+        self.steps = get_available_steps(prefix + ".", "." + self.suffix())
+        self.active_key = active_key
+
+    def _get_name_fragments(self) -> list[str]:
+        fragments = [self.prefix]
+        if self.active_key is not None:
+            fragments.append(self.active_key)
+        return fragments
+
+
+LOADERS: list[type[Loader]] = []
+
+
+def loader[T: type[Loader]](cls: T) -> T:
+    """Register a loader class. Each loader's discover() classmethod determines
+    which prefixes it claims for a given data dir."""
+    LOADERS.append(cls)
+    return cls
+
+
+def discover_loaders(data_dir: Path) -> dict[str, type[Loader]]:
+    """Poll every registered loader for the prefixes it claims in data_dir.
+    On conflict, the later-registered loader wins (with a warning), so
+    user-defined loaders can shadow built-ins."""
+    result: dict[str, type[Loader]] = {}
+    for cls in LOADERS:
+        for prefix in cls.discover_prefixes(data_dir):
+            if prefix in result:
+                warnings.warn(
+                    f"prefix '{prefix}' claimed by both {result[prefix].__name__} and {cls.__name__}; using {cls.__name__}",
+                    UserWarning,
+                    stacklevel=2,
+                )
+            result[prefix] = cls
+    return result

src/lib/data/loaders/field_bp.py

@@ -1,27 +1,33 @@
+import re
 from pathlib import Path
 
 import pscpy
 import xarray as xr
 
 from lib.config import CONFIG
 from lib.data.data_with_attrs import Field, FieldMetadata
-from lib.data.loader_registry import loader
-from lib.data.source import DataSource
+from lib.data.loader import Loader, loader
 from lib.derived_field_variables import derive_field_variable
-from lib.file_util import get_available_steps
 from lib.var_info_registry import lookup
 
+_KNOWN_PREFIXES = ("pfd", "pfd_moments", "gauss", "continuity")
+_STEP_BP_RE = re.compile(r"^(.+?)\.\d+\.bp$")
+
 
 def _get_path(prefix: str, step: int) -> Path:
     return CONFIG.data_dir / f"{prefix}.{step:09}.bp"
 
 
-@loader("pfd", "pfd_moments", "gauss", "continuity")
-class FieldLoaderBp(DataSource):
-    def __init__(self, prefix: str, active_key: str | None):
-        self.prefix = prefix
-        self.active_key = active_key
-        self.steps = get_available_steps(f"{prefix}.", ".bp")
+@loader
+class FieldLoaderBp(Loader):
+    @classmethod
+    def discover_prefixes(cls, data_dir: Path) -> list[str]:
+        present = {m.group(1) for entry in data_dir.iterdir() if (m := _STEP_BP_RE.match(entry.name))}
+        return [p for p in _KNOWN_PREFIXES if p in present]
+
+    @classmethod
+    def suffix(cls):
+        return "bp"
 
     def get_data(self) -> Field:
         ds = xr.open_mfdataset(
@@ -40,9 +46,3 @@ def get_data(self) -> Field:
             var_infos=var_info,
         )
         return Field(ds, metadata)
-
-    def _get_name_fragments(self) -> list[str]:
-        fragments = [self.prefix]
-        if self.active_key is not None:
-            fragments.append(self.active_key)
-        return fragments

src/lib/data/loaders/particle_bp.py

@@ -7,26 +7,13 @@
 
 from lib.config import CONFIG
 from lib.data.data_with_attrs import LazyList, ListMetadata
-from lib.data.loader_registry import register_loader
-from lib.data.source import DataSource
-from lib.file_util import get_available_steps
+from lib.data.loader import Loader, loader
 from lib.species import SpeciesInfo, build_species_display
 from lib.var_info_registry import lookup
 
 _DISCOVER_PARTICLE_BP_PREFIX_RE = re.compile(r"^prt\.([^.]+)\.\d+\.bp$")
 
 
-def discover_particle_bp_loaders(data_dir: pathlib.Path):
-    for entry in data_dir.iterdir():
-        # note that BP "files" are actually directories
-        m = _DISCOVER_PARTICLE_BP_PREFIX_RE.match(entry.name)
-        if m is None:
-            continue
-
-        species_key = m.group(1)
-        register_loader(f"prt.{species_key}", ParticleLoaderBp)
-
-
 def _get_path(prefix: str, step: int) -> pathlib.Path:
     return CONFIG.data_dir / f"{prefix}.{step:09}.bp"
 
@@ -58,19 +45,25 @@ def _load_step_df(path: pathlib.Path, time: float) -> dd.DataFrame:
 _SPECIES_KEY_RE = re.compile(r"^([a-zA-Z]+)([+-]*)(\d*)$")
 
 
-class ParticleLoaderBp(DataSource):
-    """ADIOS2 particle loader — one instance per prt.<species_key> prefix.
+@loader
+class ParticleLoaderBp(Loader):
+    """ADIOS2 particle loader — one instance per prt.<species_key> prefix."""
 
-    Registered dynamically by lib.parsing.parse._get_parser (not via @loader),
-    because the set of valid prefixes depends on files present in the data
-    directory at run time.
-    """
+    @classmethod
+    def discover_prefixes(cls, data_dir: pathlib.Path) -> list[str]:
+        prefixes = set()
+        for entry in data_dir.iterdir():
+            if m := _DISCOVER_PARTICLE_BP_PREFIX_RE.match(entry.name):
+                prefixes.add(f"prt.{m.group(1)}")
+        return sorted(prefixes)
+
+    @classmethod
+    def suffix(cls):
+        return "bp"
 
-    def __init__(self, prefix: str, active_key: str | None):
-        self.prefix = prefix
+    def __init__(self, prefix: str, active_key: str | None = None):
+        super().__init__(prefix, active_key)
         self.species_key = prefix.split(".", 1)[1]
-        self.active_key = active_key
-        self.steps = get_available_steps(f"{prefix}.", ".bp")
 
     def get_data(self) -> LazyList:
         step_attrs = [_read_attrs(_get_path(self.prefix, step)) for step in self.steps]
@@ -119,9 +112,3 @@ def get_data(self) -> LazyList:
             active_key=self.active_key,
             var_infos=var_infos,
         )
-
-    def _get_name_fragments(self) -> list[str]:
-        fragments = [self.prefix]
-        if self.active_key is not None:
-            fragments.append(self.active_key)
-        return fragments

src/lib/data/loaders/particle_h5.py

@@ -1,4 +1,5 @@
 import pathlib
+import re
 import typing
 import warnings
 from collections import defaultdict
@@ -9,14 +10,13 @@
 
 from lib.config import CONFIG
 from lib.data.data_with_attrs import LazyList, ListMetadata
-from lib.data.loader_registry import loader
-from lib.data.source import DataSource
-from lib.file_util import get_available_steps
+from lib.data.loader import Loader, loader
 from lib.latex import Latex
 from lib.species import SpeciesInfo, build_species_display
 from lib.var_info_registry import lookup
 
 PRT_PARTICLES_KEY = "particles/p0/1d"
+_PRT_H5_RE = re.compile(r"^prt\.\d+\.h5$")
 type SpeciesIdx = int
 type Charge = float
 type Mass = float
@@ -163,12 +163,18 @@ def _build_species_dict(qm: dict[SpeciesIdx, tuple[Charge, Mass]]) -> dict[str,
     return result
 
 
-@loader("prt")
-class ParticleLoaderH5(DataSource):
-    def __init__(self, prefix: str, active_key: str | None):
-        self.prefix = prefix
-        self.active_key = active_key
-        self.steps = get_available_steps(f"{prefix}.", ".h5")
+@loader
+class ParticleLoaderH5(Loader):
+    @classmethod
+    def discover_prefixes(cls, data_dir: pathlib.Path) -> list[str]:
+        for entry in data_dir.iterdir():
+            if _PRT_H5_RE.match(entry.name):
+                return ["prt"]
+        return []
+
+    @classmethod
+    def suffix(cls):
+        return "h5"
 
     def get_data(self) -> LazyList:
         species_dict = _build_species_dict(_discover_species_qm(self.prefix, self.steps))
@@ -206,9 +212,3 @@ def get_data(self) -> LazyList:
             var_infos=var_infos,
             subject=Latex(r"\text{Particles}"),
         )
-
-    def _get_name_fragments(self) -> list[str]:
-        fragments = [self.prefix]
-        if self.active_key is not None:
-            fragments.append(self.active_key)
-        return fragments

src/lib/parsing/args.py

@@ -3,14 +3,15 @@
 
 from lib.data.adaptor import Adaptor
 from lib.data.compile import compile_source
-from lib.data.loader_registry import LOADERS
+from lib.data.data_source import DataSource
 from lib.plotting.get_plot import get_plot
 from lib.plotting.hook import Hook
 from lib.plotting.plot import Plot
 
 
 class Args(argparse.Namespace):
     prefix: str
+    loader: DataSource
     variable: str | None
     adaptors: list[Adaptor]
     hooks: list[Hook]
@@ -19,9 +20,7 @@ class Args(argparse.Namespace):
     save_format: str | None
 
     def get_animation(self) -> Plot:
-        loader = LOADERS[self.prefix](self.prefix, self.variable)
-
-        source = compile_source(loader, self.adaptors)
+        source = compile_source(self.loader, self.adaptors)
         data = source.get_data()
 
         plot = get_plot(data)

src/lib/parsing/parse.py

@@ -1,19 +1,17 @@
 import argparse
 from pathlib import Path
+from typing import Iterable
 
 from lib.config import CONFIG
-from lib.data.loader_registry import LOADERS
-from lib.data.loaders.particle_bp import discover_particle_bp_loaders
+from lib.data.loader import discover_loaders
 from lib.parsing.args import Args
 from lib.parsing.args_registry import CUSTOM_ARGS
 
 
-def _get_parser() -> argparse.ArgumentParser:
-    # FIXME: this is a hack. Shouldn't modify a global for this.
-    discover_particle_bp_loaders(CONFIG.data_dir)
+def _get_parser(prefixes: Iterable[str]) -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(prog="psc-plot")
 
-    parser.add_argument("prefix", choices=LOADERS.keys(), help="data file prefix")
+    parser.add_argument("prefix", choices=prefixes, help="data file prefix (auto-discovered from the data directory)")
     parser.add_argument("variable", nargs="?", default=None, help="field variable to work with")
     parser.add_argument(
         "-s",
@@ -40,6 +38,9 @@ def _get_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def get_parsed_args() -> Args:
-    parser = _get_parser()
-    return parser.parse_args(namespace=Args())
+def get_parsed_args(args_list: list[str] | None = None) -> Args:
+    prefix_to_loader = discover_loaders(CONFIG.data_dir)
+    parser = _get_parser(prefix_to_loader.keys())
+    args = parser.parse_args(args_list, namespace=Args())
+    args.loader = prefix_to_loader[args.prefix](args.prefix, active_key=args.variable)
+    return args

tests/conftest.py

@@ -15,8 +15,7 @@
 import pytest
 
 from lib.config import CONFIG
-from lib.parsing.args import Args
-from lib.parsing.parse import _get_parser
+from lib.parsing.parse import get_parsed_args
 from lib.plotting.plot import SaveFormat
 
 
@@ -27,8 +26,7 @@ def make_plot(args_list: list[str], data_dir: str | None = None):
         CONFIG.data_dir = _DATA_DIR / data_dir
 
     try:
-        parser = _get_parser()
-        args = parser.parse_args(args_list, namespace=Args())
+        args = get_parsed_args(args_list)
         plot = args.get_animation()
         plot._initialize()
         return plot.fig
@@ -44,8 +42,7 @@ def make_save(args_list: list[str], save_dir: Path, format: SaveFormat, data_dir
         CONFIG.data_dir = _DATA_DIR / data_dir
 
     try:
-        parser = _get_parser()
-        args = parser.parse_args(args_list, namespace=Args())
+        args = get_parsed_args(args_list)
         plot = args.get_animation()
         save_dir.mkdir(exist_ok=True)
         plot.save(save_dir, format=format)