anemoi-datasets 0.5.26__py3-none-any.whl → 0.5.28__py3-none-any.whl

anemoi/datasets/create/__init__.py

@@ -16,8 +16,6 @@ import uuid
 import warnings
 from functools import cached_property
 from typing import Any
-from typing import Optional
-from typing import Union
 
 import cftime
 import numpy as np
@@ -102,8 +100,8 @@ def json_tidy(o: Any) -> Any:
 
 def build_statistics_dates(
     dates: list[datetime.datetime],
-    start: Optional[datetime.datetime],
-    end: Optional[datetime.datetime],
+    start: datetime.datetime | None,
+    end: datetime.datetime | None,
 ) -> tuple[str, str]:
     """Compute the start and end dates for the statistics.
 
@@ -359,7 +357,7 @@ class Actor:  # TODO: rename to Creator
 
     dataset_class = WritableDataset
 
-    def __init__(self, path: str, cache: Optional[str] = None):
+    def __init__(self, path: str, cache: str | None = None):
         """Initialize an Actor instance.
 
         Parameters
@@ -577,10 +575,10 @@ class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         check_name: bool = False,
         overwrite: bool = False,
         use_threads: bool = False,
-        statistics_temp_dir: Optional[str] = None,
+        statistics_temp_dir: str | None = None,
         progress: Any = None,
         test: bool = False,
-        cache: Optional[str] = None,
+        cache: str | None = None,
         **kwargs: Any,
     ):
         """Initialize an Init instance.
@@ -809,11 +807,11 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
     def __init__(
         self,
         path: str,
-        parts: Optional[str] = None,
+        parts: str | None = None,
         use_threads: bool = False,
-        statistics_temp_dir: Optional[str] = None,
+        statistics_temp_dir: str | None = None,
         progress: Any = None,
-        cache: Optional[str] = None,
+        cache: str | None = None,
         **kwargs: Any,
     ):
         """Initialize a Load instance.
@@ -867,7 +865,7 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
             # assert isinstance(group[0], datetime.datetime), type(group[0])
             LOG.debug(f"Building data for group {igroup}/{self.n_groups}")
 
-            result = self.input.select(group_of_dates=group)
+            result = self.input.select(argument=group)
             assert result.group_of_dates == group, (len(result.group_of_dates), len(group), group)
 
             # There are several groups.
@@ -907,8 +905,8 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
                 print("Requested dates", compress_dates(dates))
                 print("Cube dates", compress_dates(dates_in_data))
 
-                a = set(as_datetime(_) for _ in dates)
-                b = set(as_datetime(_) for _ in dates_in_data)
+                a = {as_datetime(_) for _ in dates}
+                b = {as_datetime(_) for _ in dates_in_data}
 
                 print("Missing dates", compress_dates(a - b))
                 print("Extra dates", compress_dates(b - a))
@@ -958,7 +956,7 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         array.flush()
         LOG.info("Flushed data array")
 
-    def _get_allow_nans(self) -> Union[bool, list]:
+    def _get_allow_nans(self) -> bool | list:
         """Get the allow_nans configuration.
 
         Returns
@@ -991,7 +989,7 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         total = cube.count(reading_chunks)
         LOG.debug(f"Loading datacube: {cube}")
 
-        def position(x: Any) -> Optional[int]:
+        def position(x: Any) -> int | None:
             if isinstance(x, str) and "/" in x:
                 x = x.split("/")
                 return int(x[0])
@@ -1038,7 +1036,7 @@ class Cleanup(Actor, HasRegistryMixin, HasStatisticTempMixin):
     def __init__(
         self,
         path: str,
-        statistics_temp_dir: Optional[str] = None,
+        statistics_temp_dir: str | None = None,
         delta: list = [],
         use_threads: bool = False,
         **kwargs: Any,
@@ -1217,19 +1215,19 @@ class _InitAdditions(Actor, HasRegistryMixin, AdditionsMixin):
         LOG.info(f"Cleaned temporary storage {self.tmp_storage_path}")
 
 
-class _RunAdditions(Actor, HasRegistryMixin, AdditionsMixin):
+class _LoadAdditions(Actor, HasRegistryMixin, AdditionsMixin):
     """A class to run dataset additions."""
 
     def __init__(
         self,
         path: str,
         delta: str,
-        parts: Optional[str] = None,
+        parts: str | None = None,
         use_threads: bool = False,
         progress: Any = None,
         **kwargs: Any,
     ):
-        """Initialize a _RunAdditions instance.
+        """Initialize a _LoadAdditions instance.
 
         Parameters
         ----------
@@ -1469,7 +1467,7 @@ def multi_addition(cls: type) -> type:
 
 
 InitAdditions = multi_addition(_InitAdditions)
-RunAdditions = multi_addition(_RunAdditions)
+LoadAdditions = multi_addition(_LoadAdditions)
 FinaliseAdditions = multi_addition(_FinaliseAdditions)
 
 
@@ -1480,7 +1478,7 @@ class Statistics(Actor, HasStatisticTempMixin, HasRegistryMixin):
         self,
         path: str,
         use_threads: bool = False,
-        statistics_temp_dir: Optional[str] = None,
+        statistics_temp_dir: str | None = None,
         progress: Any = None,
         **kwargs: Any,
     ):
@@ -1539,7 +1537,7 @@ class Statistics(Actor, HasStatisticTempMixin, HasRegistryMixin):
         LOG.info(f"Wrote statistics in {self.path}")
 
     @cached_property
-    def allow_nans(self) -> Union[bool, list]:
+    def allow_nans(self) -> bool | list:
         """Check if NaNs are allowed."""
         import zarr
 
@@ -1581,7 +1579,7 @@ def chain(tasks: list) -> type:
     return Chain
 
 
-def creator_factory(name: str, trace: Optional[str] = None, **kwargs: Any) -> Any:
+def creator_factory(name: str, trace: str | None = None, **kwargs: Any) -> Any:
     """Create a dataset creator.
 
     Parameters
@@ -1612,10 +1610,50 @@ def creator_factory(name: str, trace: Optional[str] = None, **kwargs: Any) -> An
         cleanup=Cleanup,
         verify=Verify,
         init_additions=InitAdditions,
-        load_additions=RunAdditions,
-        run_additions=RunAdditions,
+        load_additions=LoadAdditions,
         finalise_additions=chain([FinaliseAdditions, Size]),
-        additions=chain([InitAdditions, RunAdditions, FinaliseAdditions, Size, Cleanup]),
+        additions=chain([InitAdditions, LoadAdditions, FinaliseAdditions, Size, Cleanup]),
     )[name]
     LOG.debug(f"Creating {cls.__name__} with {kwargs}")
     return cls(**kwargs)
+
+
+def validate_config(config: Any) -> None:
+
+    import json
+
+    import jsonschema
+
+    def _tidy(d):
+        if isinstance(d, dict):
+            return {k: _tidy(v) for k, v in d.items()}
+
+        if isinstance(d, list):
+            return [_tidy(v) for v in d if v is not None]
+
+        # jsonschema does not support datetime.date
+        if isinstance(d, datetime.datetime):
+            return d.isoformat()
+
+        if isinstance(d, datetime.date):
+            return d.isoformat()
+
+        return d
+
+    # https://json-schema.org
+
+    with open(
+        os.path.join(
+            os.path.dirname(os.path.dirname(__file__)),
+            "schemas",
+            "recipe.json",
+        )
+    ) as f:
+        schema = json.load(f)
+
+    try:
+        jsonschema.validate(instance=_tidy(config), schema=schema)
+    except jsonschema.exceptions.ValidationError as e:
+        LOG.error("❌ Config validation failed (jsonschema):")
+        LOG.error(e.message)
+        raise

anemoi/datasets/create/check.py

@@ -12,10 +12,8 @@ import datetime
 import logging
 import re
 import warnings
+from collections.abc import Callable
 from typing import Any
-from typing import Callable
-from typing import Optional
-from typing import Union
 
 import numpy as np
 from anemoi.utils.config import load_config
@@ -31,10 +29,10 @@ class DatasetName:
     def __init__(
         self,
         name: str,
-        resolution: Optional[str] = None,
-        start_date: Optional[datetime.date] = None,
-        end_date: Optional[datetime.date] = None,
-        frequency: Optional[datetime.timedelta] = None,
+        resolution: str | None = None,
+        start_date: datetime.date | None = None,
+        end_date: datetime.date | None = None,
+        frequency: datetime.timedelta | None = None,
     ):
         """Initialize a DatasetName instance.
 
@@ -146,7 +144,7 @@ class DatasetName:
                 "https://anemoi-registry.readthedocs.io/en/latest/naming-conventions.html"
             )
 
-    def check_resolution(self, resolution: Optional[str]) -> None:
+    def check_resolution(self, resolution: str | None) -> None:
         """Check if the resolution matches the expected format.
 
         Parameters
@@ -175,7 +173,7 @@ class DatasetName:
             if not c.isalnum() and c not in "-":
                 self.messages.append(f"the {self.name} should only contain alphanumeric characters and '-'.")
 
-    def check_frequency(self, frequency: Optional[datetime.timedelta]) -> None:
+    def check_frequency(self, frequency: datetime.timedelta | None) -> None:
         """Check if the frequency matches the expected format.
 
         Parameters
@@ -189,7 +187,7 @@ class DatasetName:
         self._check_missing("frequency", frequency_str)
         self._check_mismatch("frequency", frequency_str)
 
-    def check_start_date(self, start_date: Optional[datetime.date]) -> None:
+    def check_start_date(self, start_date: datetime.date | None) -> None:
         """Check if the start date matches the expected format.
 
         Parameters
@@ -203,7 +201,7 @@ class DatasetName:
         self._check_missing("start_date", start_date_str)
         self._check_mismatch("start_date", start_date_str)
 
-    def check_end_date(self, end_date: Optional[datetime.date]) -> None:
+    def check_end_date(self, end_date: datetime.date | None) -> None:
         """Check if the end date matches the expected format.
 
         Parameters
@@ -251,7 +249,7 @@ class StatisticsValueError(ValueError):
 
 
 def check_data_values(
-    arr: NDArray[Any], *, name: str, log: list = [], allow_nans: Union[bool, list, set, tuple, dict] = False
+    arr: NDArray[Any], *, name: str, log: list = [], allow_nans: bool | list | set | tuple | dict = False
 ) -> None:
     """Check the values in the data array for validity.
 

anemoi/datasets/create/chunks.py

@@ -9,7 +9,6 @@
 
 import logging
 import warnings
-from typing import Union
 
 LOG = logging.getLogger(__name__)
 
@@ -27,7 +26,7 @@ class ChunkFilter:
         The chunks that are allowed to be processed.
     """
 
-    def __init__(self, *, parts: Union[str, list], total: int):
+    def __init__(self, *, parts: str | list, total: int):
         """Initializes the ChunkFilter with the given parts and total number of chunks.
 
         Parameters

anemoi/datasets/create/config.py

@@ -12,8 +12,6 @@ import logging
 import os
 from copy import deepcopy
 from typing import Any
-from typing import Optional
-from typing import Union
 
 import yaml
 from anemoi.utils.config import DotDict
@@ -25,7 +23,7 @@ from anemoi.datasets.dates.groups import Groups
 LOG = logging.getLogger(__name__)
 
 
-def _get_first_key_if_dict(x: Union[str, dict]) -> str:
+def _get_first_key_if_dict(x: str | dict) -> str:
     """Returns the first key if the input is a dictionary, otherwise returns the input string.
 
     Parameters
@@ -97,7 +95,7 @@ def check_dict_value_and_set(dic: dict, key: str, value: Any) -> None:
     dic[key] = value
 
 
-def resolve_includes(config: Union[dict, list]) -> Union[dict, list]:
+def resolve_includes(config: dict | list) -> dict | list:
     """Resolves '<<' includes in a configuration dictionary or list.
 
     Parameters
@@ -123,7 +121,7 @@ def resolve_includes(config: Union[dict, list]) -> Union[dict, list]:
 class Config(DotDict):
     """Configuration class that extends DotDict to handle configuration loading and processing."""
 
-    def __init__(self, config: Optional[Union[str, dict]] = None, **kwargs):
+    def __init__(self, config: str | dict | None = None, **kwargs):
         """Initializes the Config object.
 
         Parameters
@@ -134,7 +132,6 @@ class Config(DotDict):
             Additional keyword arguments to update the configuration.
         """
         if isinstance(config, str):
-            self.config_path = os.path.realpath(config)
             config = load_any_dict_format(config)
         else:
             config = deepcopy(config if config is not None else {})
@@ -282,6 +279,8 @@ class LoadersConfig(Config):
 
         self.output.order_by = normalize_order_by(self.output.order_by)
 
+        self.setdefault("dates", Config())
+
         self.dates["group_by"] = self.build.group_by
 
         ###########

anemoi/datasets/create/input/__init__.py

@@ -1,4 +1,4 @@
-# (C) Copyright 2024 Anemoi contributors.
+# (C) Copyright 2024-2025 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
@@ -7,100 +7,79 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-import logging
 from copy import deepcopy
+from functools import cached_property
+from typing import TYPE_CHECKING
 from typing import Any
-from typing import Union
 
-from anemoi.datasets.dates.groups import GroupOfDates
+from anemoi.datasets.create.input.context.field import FieldContext
 
-from .trace import trace_select
-
-LOG = logging.getLogger(__name__)
-
-
-class Context:
-    """Context for building input data."""
-
-    pass
+if TYPE_CHECKING:
+    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: Union[dict, list], **kwargs: Any) -> None:
+    def __init__(self, config: dict, data_sources: dict | list, **kwargs: Any) -> None:
         """Initialize the InputBuilder.
 
         Parameters
         ----------
         config : dict
             Configuration dictionary.
-        data_sources : Union[dict, list]
+        data_sources : dict
             Data sources.
         **kwargs : Any
             Additional keyword arguments.
         """
         self.kwargs = kwargs
+        self.config = deepcopy(config)
+        self.data_sources = deepcopy(dict(data_sources=data_sources))
 
-        config = deepcopy(config)
-        if data_sources:
-            config = dict(
-                data_sources=dict(
-                    sources=data_sources,
-                    input=config,
-                )
-            )
-        self.config = config
-        self.action_path = ["input"]
-
-    @trace_select
-    def select(self, group_of_dates: GroupOfDates) -> Any:
-        """Select data based on the group of dates.
-
-        Parameters
-        ----------
-        group_of_dates : GroupOfDates
-            Group of dates to select data for.
-
-        Returns
-        -------
-        Any
-            Selected data.
-        """
-        from .action import ActionContext
+    @cached_property
+    def action(self) -> "Recipe":
+        """Returns the action object based on the configuration."""
+        from .action import Recipe
         from .action import action_factory
 
-        """This changes the context."""
-        context = ActionContext(**self.kwargs)
-        action = action_factory(self.config, context, self.action_path)
-        return action.select(group_of_dates)
-
-    def __repr__(self) -> str:
-        """Return a string representation of the InputBuilder.
-
-        Returns
-        -------
-        str
-            String representation.
-        """
-        from .action import ActionContext
-        from .action import action_factory
+        sources = action_factory(self.data_sources, "data_sources")
+        input = action_factory(self.config, "input")
 
-        context = ActionContext(**self.kwargs)
-        a = action_factory(self.config, context, self.action_path)
-        return repr(a)
+        return Recipe(input, sources)
 
-    def _trace_select(self, group_of_dates: GroupOfDates) -> str:
-        """Trace the select operation.
+    def select(self, argument) -> Any:
+        """Select data based on the group of dates.
 
         Parameters
         ----------
-        group_of_dates : GroupOfDates
+        argument : GroupOfDates
             Group of dates to select data for.
 
         Returns
         -------
-        str
-            Trace string.
+        Any
+            Selected data.
         """
-        return f"InputBuilder({group_of_dates})"
+        context = FieldContext(argument, **self.kwargs)
+        return context.create_result(self.action(context, argument))
+
+
+def build_input(config: dict, data_sources: dict | list, **kwargs: Any) -> InputBuilder:
+    """Build an InputBuilder instance.
+
+    Parameters
+    ----------
+    config : dict
+        Configuration dictionary.
+    data_sources : Union[dict, list]
+        Data sources.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    InputBuilder
+        An instance of InputBuilder.
+    """
+    return InputBuilder(config, data_sources, **kwargs)