feat: pydantic recipes (#575)

## Description

Use Pydandic to load the recipe files, including models that my be
defined in sources, filters and other plugins.

To add a pydantic mode to a filter:


```python

class Schema(BaseModel):
   fermi: str = "paradox"
   kardashev_scale: int = 3
   
   
@filter_registry.register("the-great-filter")
class TheGreatFilter(Filter):

      schema = Schema
      
      def __init__(self, fermi, kardashev_scale):
         ....
         
``` 

## What problem does this change solve?
<!-- Describe if it's a bugfix, new feature, doc update, or breaking
change -->

## What issue or task does this change relate to?
<!-- link to Issue Number -->

##  Additional notes ##
<!-- Include any additional information, caveats, or considerations that
the reviewer should be aware of. -->

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

---------

Co-authored-by: Oskar Weser <oskar.weser@gmail.com>

Author: b8raoult

src/anemoi/datasets/commands/inspect.py

@@ -740,13 +740,18 @@ class Version0_14(Version0_13):
     pass
 
 
+class Version0_15(Version0_14):
+    pass
+
+
 VERSIONS = {
     "0.0.0": NoVersion,
     "0.4.0": Version0_4,
     "0.6.0": Version0_6,
     "0.12.0": Version0_12,
     "0.13.0": Version0_13,
     "0.14.0": Version0_14,
+    "0.15.0": Version0_15,
 }
 
 

src/anemoi/datasets/commands/recipe/dump.py

@@ -41,7 +41,7 @@ def dump_recipe(config: dict, dumper=None) -> str:
     recipe = Recipe(**config)
     input = InputBuilder(
         recipe.input,
-        data_sources=recipe.data_sources or {},
+        data_sources=recipe.data_sources,
     )
     if dumper is None:
         dumper = Dumper()

src/anemoi/datasets/create/creator.py

@@ -21,6 +21,7 @@
 import tqdm
 import zarr
 from anemoi.utils.humanize import bytes_to_human
+from anemoi.utils.sanitise import sanitise
 
 from anemoi.datasets.create.dataset import Dataset
 from anemoi.datasets.create.input.builder import InputBuilder
@@ -32,7 +33,7 @@
 
 LOG = logging.getLogger(__name__)
 
-VERSION = "0.14"
+VERSION = "0.15"
 
 LOG = logging.getLogger(__name__)
 
@@ -177,9 +178,10 @@ def fill_metadata(self, metadata: dict) -> None:
         model_dump = json.loads(model_dump)
         model_dump = self.recipe.strip_unknown_keys(model_dump)
 
-        # TODO: make it an option
-        # recipe = sanitise(model_dump)
-        recipe = model_dump
+        if self.recipe.output.sanitise:
+            recipe = sanitise(model_dump)
+        else:
+            recipe = model_dump
 
         # Remove stuff added by prepml
         allow_keys = set(model_dump.keys())
@@ -348,7 +350,7 @@ def task_finalise_additions(self) -> None:
     @cached_property
     def groups(self) -> Groups:
         """Return the date groups for the dataset."""
-        return Groups(**self.recipe.dates, group_by=self.recipe.build.group_by)
+        return Groups(self.recipe.dates, group_by=self.recipe.build.group_by)
 
     @cached_property
     def minimal_input(self) -> Any:
@@ -362,7 +364,7 @@ def input(self) -> InputBuilder:
 
         return InputBuilder(
             self.recipe.input,
-            data_sources=self.recipe.data_sources or {},
+            data_sources=self.recipe.data_sources,
         )
 
     @cached_property

src/anemoi/datasets/create/input/action.py

@@ -10,8 +10,9 @@
 import logging
 from abc import ABC
 from abc import abstractmethod
+from functools import cache
 
-from anemoi.datasets.dates import DatesProvider
+from anemoi.datasets.create.recipe.dates import StartEndDates
 
 LOG = logging.getLogger(__name__)
 
@@ -79,7 +80,7 @@ def __init__(self, config, *path):
         for i, item in enumerate(config):
 
             dates = item["dates"]
-            filtering_dates = DatesProvider.from_config(**dates)
+            filtering_dates = StartEndDates(**dates)
             action = action_factory({k: v for k, v in item.items() if k != "dates"}, *self.path, str(i))
             self.choices.append((filtering_dates, action))
 
@@ -308,48 +309,55 @@ def __call__(self, context, argument):
         return self.input(context, argument)
 
 
-KLASS = {
-    "concat": Concat,
-    "join": Join,
-    "pipe": Pipe,
-    "data-sources": DataSources,
-}
+@cache
+def _action_factories():
 
-LEN_KLASS = len(KLASS)
+    factories = {
+        "concat": Concat,
+        "join": Join,
+        "pipe": Pipe,
+        "data-sources": DataSources,
+    }
 
+    # Load pluggins
+    from anemoi.transform.filters import filter_registry as transform_filter_registry
+    from anemoi.transform.sources import source_registry as transform_source_registry
 
-def make(key, config, *path):
+    from anemoi.datasets.create.sources import source_registry as dataset_source_registry
+
+    # Register sources, local first
+    for name in dataset_source_registry.registered:
+        if name not in factories:
+            factories[name.replace("_", "-")] = new_source(name, DatasetSourceMixin)
 
-    if LEN_KLASS == len(KLASS):
+    for name in transform_source_registry.registered:
+        if name not in factories:
+            factories[name.replace("_", "-")] = new_source(name, TransformSourceMixin)
 
-        # Load pluggins
-        from anemoi.transform.filters import filter_registry as transform_filter_registry
-        from anemoi.transform.sources import source_registry as transform_source_registry
+    # Register filters
+    for name in transform_filter_registry.registered:
+        if name not in factories:
+            factories[name.replace("_", "-")] = new_filter(name, TransformFilterMixin)
 
-        from anemoi.datasets.create.sources import source_registry as dataset_source_registry
+    return factories
 
-        # Register sources, local first
-        for name in dataset_source_registry.registered:
-            if name not in KLASS:
-                KLASS[name.replace("_", "-")] = new_source(name, DatasetSourceMixin)
 
-        for name in transform_source_registry.registered:
-            if name not in KLASS:
-                KLASS[name.replace("_", "-")] = new_source(name, TransformSourceMixin)
+def make(key, config, *path):
 
-        # Register filters
-        for name in transform_filter_registry.registered:
-            if name not in KLASS:
-                KLASS[name.replace("_", "-")] = new_filter(name, TransformFilterMixin)
+    factories = _action_factories()
 
-    return KLASS[key.replace("_", "-")](config, *path)
+    return factories[key.replace("_", "-")](config, *path)
 
 
 def action_factory(data, *path):
+    from pydantic import BaseModel
 
     assert len(path) > 0, f"Path must contain at least one element {path}"
     assert path[0] in ("input", "data_sources")
 
+    if isinstance(data, BaseModel):
+        data = data.model_dump()
+
     assert isinstance(data, dict), f"Input data must be a dictionary, got {type(data)}"
     assert len(data) == 1, f"Input data must contain exactly one key-value pair {data} {'.'.join(x for x in path)}"
 

src/anemoi/datasets/create/input/builder.py

@@ -7,33 +7,37 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+from __future__ import annotations
+
 from copy import deepcopy
 from functools import cached_property
 from typing import TYPE_CHECKING
 from typing import Any
 
 if TYPE_CHECKING:
+    from anemoi.datasets.create.input.action import Action
     from anemoi.datasets.create.input.action import Recipe
 
 
 class InputBuilder:
     """Builder class for creating input data from configuration and data sources."""
 
-    def __init__(self, config: dict, data_sources: dict | list, **kwargs: Any) -> None:
+    def __init__(self, config: "Action" | None, data_sources: "Action" | None, **kwargs: Any) -> None:
         """Initialize the InputBuilder.
 
         Parameters
         ----------
-        config : dict
-            Configuration dictionary.
-        data_sources : dict
-            Data sources.
+        config : Action | None
+            Input part of the recipe.
+        data_sources : Action   | None
+            Data sources part of the recipe.
         **kwargs : Any
             Additional keyword arguments.
         """
         self.kwargs = kwargs
-        self.config = deepcopy(config)
-        self.data_sources = deepcopy(dict(data_sources=data_sources))
+        # self.config = deepcopy(config.model_dump() if config is not None else {"empty": {}})
+        self.config = deepcopy(config.model_dump() if config is not None else {})
+        self.data_sources = deepcopy(dict(data_sources=data_sources or {}))
 
     @cached_property
     def action(self) -> "Recipe":

src/anemoi/datasets/create/recipe/__init__.py

@@ -9,35 +9,27 @@
 
 import json
 import logging
-from typing import Annotated
 
 import yaml
-from anemoi.utils.config import DotDict
 from pydantic import BaseModel
-from pydantic import BeforeValidator
 from pydantic import ConfigDict
 from pydantic import Field
 from pydantic import model_validator
 
+from .action import Action
 from .build import Build
+from .dates import Dates
 from .output import GriddedOutput
 from .output import Output
 from .statistics import Statistics
 
 LOG = logging.getLogger(__name__)
 
 
-def validate_dotdict(v):
-    if isinstance(v, dict):
-        return DotDict(v)
-    return v
-
-
-DotDictField = Annotated[DotDict, BeforeValidator(validate_dotdict)]
-
-
 class Recipe(BaseModel):
 
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
+
     @model_validator(mode="after")
     def _post_init(self) -> "Recipe":
         # We need to call _post_init on nested BaseModel members
@@ -48,27 +40,26 @@ def _post_init(self) -> "Recipe":
                 member._post_init(self)
         return self
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
-
     description: str = "No description provided."
     licence: str = "unknown"
     attribution: str = "unknown"
 
-    dates: DotDictField
+    dates: Dates
     """The date configuration for the dataset."""
 
-    input: DotDictField
+    input: Action | None = None
     """The input data sources configuration."""
 
-    data_sources: list[DotDictField] | DotDictField | None = None
+    data_sources: dict[str, Action] | list[Action] | None = None
     """The data sources configuration."""
 
     output: Output = Field(default_factory=GriddedOutput)
     """The output configuration."""
 
     build: Build = Build()
     """The build configuration."""
-    additions: DotDictField | None = Field(
+
+    additions: dict | None = Field(
         default=None,
         deprecated="Top-level 'additions' is deprecated. Use 'statistics.tendencies' instead.",
     )
@@ -80,9 +71,48 @@ def _post_init(self) -> "Recipe":
         deprecated="Top-level 'env' is deprecated. Please use 'build.env' instead.",
     )
 
+    def only_non_defaults(self, data: dict) -> dict:
+        """Return a dictionary containing only non-default values from the recipe.
+
+        Parameters
+        ----------
+        data : dict
+            The recipe data as a dictionary.
+
+        Returns
+        -------
+        dict
+            A dictionary containing only non-default values.
+        """
+
+        defaults = Recipe(dates={"values": []}).model_dump()
+
+        def _only_non_defaults(d, default_d):
+
+            if type(d) is not type(default_d):
+                return d
+
+            if isinstance(d, dict):
+                res = d.copy()
+                for k, v in list(d.items()):
+                    if k not in default_d:
+                        del res[k]
+                        continue
+
+                    if v == default_d[k]:
+                        del res[k]
+                        continue
+
+                    res[k] = _only_non_defaults(v, default_d[k])
+                return res
+
+            return d
+
+        return _only_non_defaults(data, defaults)
+
     def strip_unknown_keys(self, data: dict) -> dict:
         assert isinstance(data, dict)
-        defaults = Recipe(dates={}, input={}).model_dump()
+        defaults = Recipe(input={"empty": {}}, dates={"values": []}).model_dump()
         return {key: data[key] for key in defaults.keys()}
 
 
@@ -99,8 +129,6 @@ def loader_recipe_from_yaml(path: str) -> dict:
     dict
         The dataset recipe.
     """
-    LOG.info(f"Loading recipe from YAML file at {path}")
-
     with open(path) as f:
         recipe_yaml = f.read()
     recipe = yaml.safe_load(recipe_yaml)
@@ -122,12 +150,12 @@ def loader_recipe_from_zarr(path: str) -> dict:
     """
     import zarr
 
-    LOG.info(f"Loading recipe from Zarr store at {path}")
-
     z = zarr.open(path, mode="r")
 
     for name in ("_recipe", "recipe"):
         if name not in z.attrs:
+            # return None
+            LOG.error(f"No '{name}' found in Zarr store at {path}")
             continue
 
         recipe = z.attrs[name]