feat: explicit pydantic schema for the pipeline config

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Changed
+- **Breaking:** the pipeline config is validated against an explicit pydantic schema (`mxalign.pipeline.config.PipelineConfig`, loaded with `load_config`) instead of being read as a free-form dict. Unknown sections or keys, and missing required keys, raise a `ValidationError` when the `Runner` is created; `Runner.config` is now the validated model. Per-dataset `dates` entries override individual fields of the top-level `dates`. @leifdenby
 - **Breaking:** file path patterns (`datasets.<name>.files` with a `dates:` section, and `alignment.save.path`) now use Python `str.format` syntax instead of earthkit-data patterns: `{reference_time:%Y%m%d%H}` instead of `{reference_time:strftime(%Y%m%d%H)}`, `{year:02d}` instead of `{year:int(%02d)}`. `{lead_time}` (integer hours) and `{valid_time}` are now supported in `files` patterns. `earthkit-data` is no longer a dependency. @michielv
 - **Breaking:** temporal alignment always runs in the pipeline; `alignment.time` is optional and only accepts `lead_time` (`reference` | `intersection` | `union`, forecast reference only). Unknown keys such as `return_as` or `method` now raise an error. `align_time` takes `reference` and an explicit `lead_time` instead of `**kwargs`. @michielv
 - Remove dataset loading and dataset validation and outsource it to [mlwp-data-loaders](https://github.com/mlwp-tools/mlwp-data-loaders) and [mlwp-data-specs](https://github.com/mlwp-tools/mlwp-data-specs) @michielv

README.md

@@ -43,7 +43,7 @@ For a more comprehensive interactive example, check out the [introductory notebo
 
 ### Configuration Contents
 
-The configuration file is divided into several main sections:
+The configuration file is divided into several main sections. It is validated on load against the schema in `mxalign.pipeline.config` (`PipelineConfig`): unknown sections or keys, and missing required ones such as `alignment.reference`, raise an error up front. Keys that are passed through to a function, such as loader, transformation, metric and save options, are accepted as-is.
 
 ```yaml
 datasets:

examples/advanced.ipynb

@@ -455,7 +455,7 @@
     }
    ],
    "source": [
-    "runner.config()"
+    "runner.config"
    ]
   },
   {

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
     "mlwp-data-specs",
     "mlwp-data-loaders",
     "pandas>=2.2",
+    "pydantic>=2",
 ]
 
 [project.scripts]

src/mxalign/pipeline/config.py

@@ -1,62 +1,136 @@
+"""Schema for the pipeline YAML config.
+
+Sections and keys not listed here are rejected, except where a model allows
+extra keys (``extra="allow"``); those are passed on as keyword arguments to
+the function the section configures (loader, transformation, metric, save).
+"""
+
+from datetime import datetime
+from typing import Literal
+
 import yaml
+from pydantic import BaseModel, ConfigDict, model_validator
 
 from .dates import Dates
 
 
-def load_yaml(fn: str) -> dict:
-    """Load a YAML file and return its contents as a dict."""
-    with open(fn, "r") as f:
-        return yaml.safe_load(f)
+class _Strict(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class _PassThrough(BaseModel):
+    """Model whose unknown keys are kept in ``model_extra`` as function kwargs."""
 
+    model_config = ConfigDict(extra="allow")
 
-class Config:
-    """Pipeline configuration loaded from a YAML file or a plain dict.
+    @property
+    def kwargs(self) -> dict:
+        return dict(self.model_extra or {})
 
-    Top-level ``dates`` entries are merged into each dataset's loader config
-    so that file-path patterns can be expanded to concrete paths.
+
+class DatesConfig(_Strict):
+    """Date range used to expand ``{reference_time}``/``{lead_time}``/``{valid_time}`` in ``files``.
+
+    Every field is optional so a dataset can override part of the top-level
+    ``dates`` section; after merging, all five must be set.
     """
 
-    def __init__(self, config: str | dict) -> None:
-        self.config = load_yaml(config) if isinstance(config, str) else config
-        if not isinstance(self.config, dict):
-            raise TypeError("config should be a dictionary.")
-        self.dates = self.config.pop("dates", None)
-        self._init_datasets()
-        print("Config initialized")
-
-    def __getitem__(self, key: str) -> dict | None:
-        config = self.config.get(key, None)
-        if config:
-            return config.copy()
-        else:
-            return config
-
-    def __call__(self) -> dict:
-        return self.config
-
-    def _init_datasets(self) -> None:
-        for key, loader in self.config["datasets"].items():
-            dates_loader = loader.pop("dates", None)
-            if self.dates:
-                if dates_loader:
-                    keys_all = list(set(self.dates.keys()).union(dates_loader.keys()))
-                    dates = {
-                        key: (
-                            dates_loader[key]
-                            if key in dates_loader.keys()
-                            else self.dates[key]
-                        )
-                        for key in keys_all
-                    }
-                else:
-                    dates = self.dates.copy()
-            else:
-                if dates_loader:
-                    dates = dates_loader.copy()
-                else:
-                    dates = None
-
-            if dates:
-                dates = Dates(**dates)
-                loader["files"] = dates.substitute(loader["files"])
-            self.config["datasets"][key] = loader
+    start: datetime | str | None = None
+    end: datetime | str | None = None
+    period: str | None = None
+    range: str | None = None
+    step: str | None = None
+
+
+class DatasetConfig(_PassThrough):
+    """One entry of ``datasets``; extra keys go to ``load_and_validate_dataset``."""
+
+    loader: str
+    files: str | list[str]
+    variables: str | list[str] | None = None
+    grid_mapping: str | dict | None = None
+    dates: DatesConfig | None = None
+
+
+class TransformationConfig(_PassThrough):
+    """One entry of ``transformations`` (keyed by transformation name); extra keys go to the transformation."""
+
+    datasets: list[str] | None = None
+
+
+class TimeAlignmentConfig(_Strict):
+    lead_time: str | list[str] | None = None
+
+
+class SpaceAlignmentConfig(_Strict):
+    """Interpolator/regridder kwargs, chosen by the kind of space alignment needed."""
+
+    interpolation: dict | None = None
+    regrid: dict | None = None
+
+
+class SaveConfig(_PassThrough):
+    """Output for ``ds.<method>(path, **kwargs)``; extra keys go to the save method."""
+
+    method: str
+    path: str
+
+
+class AlignmentSaveConfig(SaveConfig):
+    datasets: Literal["all", "merge"] = "all"
+
+
+class AlignmentConfig(_Strict):
+    reference: str
+    broadcast_nans: bool = True
+    time: TimeAlignmentConfig | None = None
+    space: SpaceAlignmentConfig | None = None
+    save: AlignmentSaveConfig | None = None
+
+
+class MetricConfig(_PassThrough):
+    """One entry of ``verification.metrics``; extra keys go to the metric function."""
+
+    function: str
+    inputs: dict[str, Literal["forecast", "reference"]]
+
+
+class VerificationConfig(_Strict):
+    reference: str
+    metrics: dict[str, MetricConfig] | None = None
+    save: SaveConfig | None = None
+
+
+class PipelineConfig(_Strict):
+    dates: DatesConfig | None = None
+    datasets: dict[str, DatasetConfig]
+    transformations: dict[str, TransformationConfig] | None = None
+    alignment: AlignmentConfig
+    verification: VerificationConfig | None = None
+
+    @model_validator(mode="after")
+    def _expand_files(self) -> "PipelineConfig":
+        """Expand each dataset's ``files`` pattern(s) over its (merged) ``dates``."""
+        top = self.dates.model_dump(exclude_none=True) if self.dates else {}
+        for name, ds in self.datasets.items():
+            own = ds.dates.model_dump(exclude_none=True) if ds.dates else {}
+            dates = {**top, **own}
+            if not dates:
+                continue
+            missing = set(DatesConfig.model_fields) - dates.keys()
+            if missing:
+                raise ValueError(
+                    f"datasets.{name}: dates is missing {', '.join(sorted(missing))}"
+                )
+            patterns = [ds.files] if isinstance(ds.files, str) else ds.files
+            dates = Dates(**dates)
+            ds.files = [path for p in patterns for path in dates.substitute(p)]
+        return self
+
+
+def load_config(config: str | dict) -> PipelineConfig:
+    """Validate a pipeline config given as a YAML file path or a dict."""
+    if isinstance(config, str):
+        with open(config, "r") as f:
+            config = yaml.safe_load(f)
+    return PipelineConfig.model_validate(config)