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

anemoi/datasets/data/subset.py

@@ -10,12 +10,9 @@
 
 import datetime
 import logging
+from collections.abc import Sequence
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Sequence
-from typing import Set
 from typing import Union
 
 import numpy as np
@@ -95,7 +92,7 @@ def _end(a: int, b: int, dates: NDArray[np.datetime64]) -> int:
         return b
 
 
-def _combine_reasons(reason1: Dict[str, Any], reason2: Dict[str, Any], dates: NDArray[np.datetime64]) -> Dict[str, Any]:
+def _combine_reasons(reason1: dict[str, Any], reason2: dict[str, Any], dates: NDArray[np.datetime64]) -> dict[str, Any]:
     """Combine two reason dictionaries.
 
     Parameters:
@@ -126,7 +123,7 @@ class Subset(Forwards):
     reason (Dict[str, Any]): Dictionary of reasons.
     """
 
-    def __init__(self, dataset: Union[Dataset, "Subset"], indices: Sequence[int], reason: Dict[str, Any]) -> None:
+    def __init__(self, dataset: Union[Dataset, "Subset"], indices: Sequence[int], reason: dict[str, Any]) -> None:
         """Initialize the Subset.
 
         Parameters:
@@ -140,8 +137,8 @@ class Subset(Forwards):
             dataset = dataset.dataset
 
         self.dataset: Dataset = dataset
-        self.indices: List[int] = list(indices)
-        self.reason: Dict[str, Any] = {k: v for k, v in reason.items() if v is not None}
+        self.indices: list[int] = list(indices)
+        self.reason: dict[str, Any] = {k: v for k, v in reason.items() if v is not None}
 
         # Forward other properties to the super dataset
         super().__init__(dataset)
@@ -274,10 +271,10 @@ class Subset(Forwards):
         return f"Subset({self.dataset},{self.dates[0]}...{self.dates[-1]}/{self.frequency})"
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the missing indices of the subset."""
         missing = self.dataset.missing
-        result: Set[int] = set()
+        result: set[int] = set()
         for j, i in enumerate(self.indices):
             if i in missing:
                 result.add(j)
@@ -291,7 +288,7 @@ class Subset(Forwards):
         """
         return Node(self, [self.dataset.tree()], **self.reason)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the forwards subclass.
 
         Returns:

anemoi/datasets/data/unchecked.py

@@ -10,14 +10,10 @@
 
 import datetime
 import logging
+from collections.abc import Callable
 from functools import cached_property
 from functools import wraps
 from typing import Any
-from typing import Callable
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Set
 
 import numpy as np
 from numpy.typing import NDArray
@@ -163,13 +159,13 @@ class Unchecked(Combined):
 
     @check("check_same_variables")
     @property
-    def name_to_index(self) -> Dict[str, int]:
+    def name_to_index(self) -> dict[str, int]:
         """Get the mapping of variable names to their indices."""
         raise NotImplementedError()
 
     @check("check_same_variables")
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Get the list of variables in the dataset."""
         raise NotImplementedError()
 
@@ -181,12 +177,12 @@ class Unchecked(Combined):
 
     @check("check_same_variables")
     @property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Get the statistics of the dataset."""
         raise NotImplementedError()
 
     @check("check_same_variables")
-    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+    def statistics_tendencies(self, delta: datetime.timedelta | None = None) -> dict[str, NDArray[Any]]:
         """Get the statistics tendencies of the dataset.
 
         Parameters
@@ -207,9 +203,9 @@ class Unchecked(Combined):
         raise NotImplementedError()
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the missing data indices."""
-        result: Set[int] = set()
+        result: set[int] = set()
         for d in self.datasets:
             result = result | d.missing
         return result

anemoi/datasets/data/xy.py

@@ -11,10 +11,6 @@
 import logging
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Set
-from typing import Tuple
 
 from .dataset import Dataset
 from .dataset import FullIndex
@@ -29,7 +25,7 @@ LOG = logging.getLogger(__name__)
 class ZipBase(Combined):
     """Base class for handling zipped datasets."""
 
-    def __init__(self, datasets: List[Any], check_compatibility: bool = True) -> None:
+    def __init__(self, datasets: list[Any], check_compatibility: bool = True) -> None:
         """Initialize ZipBase with a list of datasets.
 
         Parameters
@@ -58,7 +54,7 @@ class ZipBase(Combined):
         new_parents = [parent.clone(ds) for ds in self.datasets]
         return self.clone(new_parents)
 
-    def clone(self, datasets: List[Any]) -> "ZipBase":
+    def clone(self, datasets: list[Any]) -> "ZipBase":
         """Clone the ZipBase with new datasets.
 
         Parameters
@@ -81,7 +77,11 @@ class ZipBase(Combined):
         Node
             Tree representation of the datasets.
         """
-        return Node(self, [d.tree() for d in self.datasets], check_compatibility=self._check_compatibility)
+        return Node(
+            self,
+            [d.tree() for d in self.datasets],
+            check_compatibility=self._check_compatibility,
+        )
 
     def __len__(self) -> int:
         """Get the length of the smallest dataset.
@@ -93,7 +93,7 @@ class ZipBase(Combined):
         """
         return min(len(d) for d in self.datasets)
 
-    def __getitem__(self, n: FullIndex) -> Tuple[Any, ...]:
+    def __getitem__(self, n: FullIndex) -> tuple[Any, ...]:
         """Get the item at the specified index from all datasets.
 
         Parameters
@@ -145,55 +145,55 @@ class ZipBase(Combined):
         pass
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the set of missing indices from all datasets."""
-        result: Set[int] = set()
+        result: set[int] = set()
         for d in self.datasets:
             result = result | d.missing
         return result
 
     @property
-    def shape(self) -> Tuple[Any, ...]:
+    def shape(self) -> tuple[Any, ...]:
         """Get the shape of all datasets."""
         return tuple(d.shape for d in self.datasets)
 
     @property
-    def field_shape(self) -> Tuple[Any, ...]:
+    def field_shape(self) -> tuple[Any, ...]:
         """Get the field shape of all datasets."""
         return tuple(d.shape for d in self.datasets)
 
     @property
-    def latitudes(self) -> Tuple[Any, ...]:
+    def latitudes(self) -> tuple[Any, ...]:
         """Get the latitudes of all datasets."""
         return tuple(d.latitudes for d in self.datasets)
 
     @property
-    def longitudes(self) -> Tuple[Any, ...]:
+    def longitudes(self) -> tuple[Any, ...]:
         """Get the longitudes of all datasets."""
         return tuple(d.longitudes for d in self.datasets)
 
     @property
-    def dtype(self) -> Tuple[Any, ...]:
+    def dtype(self) -> tuple[Any, ...]:
         """Get the data types of all datasets."""
         return tuple(d.dtype for d in self.datasets)
 
     @property
-    def grids(self) -> Tuple[Any, ...]:
+    def grids(self) -> tuple[Any, ...]:
         """Get the grids of all datasets."""
         return tuple(d.grids for d in self.datasets)
 
     @property
-    def statistics(self) -> Tuple[Any, ...]:
+    def statistics(self) -> tuple[Any, ...]:
         """Get the statistics of all datasets."""
         return tuple(d.statistics for d in self.datasets)
 
     @property
-    def resolution(self) -> Tuple[Any, ...]:
+    def resolution(self) -> tuple[Any, ...]:
         """Get the resolution of all datasets."""
         return tuple(d.resolution for d in self.datasets)
 
     @property
-    def name_to_index(self) -> Tuple[Any, ...]:
+    def name_to_index(self) -> tuple[Any, ...]:
         """Get the name to index mapping of all datasets."""
         return tuple(d.name_to_index for d in self.datasets)
 
@@ -210,6 +210,10 @@ class ZipBase(Combined):
         if self._check_compatibility:
             super().check_compatibility(d1, d2)
 
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
+        """Get the metadata specific to the subclass."""
+        return {}
+
 
 class Zip(ZipBase):
     """Class for handling zipped datasets."""
@@ -223,7 +227,7 @@ class XY(ZipBase):
     pass
 
 
-def xy_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> XY:
+def xy_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> XY:
     """Factory function to create an XY instance.
 
     Parameters
@@ -256,7 +260,7 @@ def xy_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> XY:
     return XY(datasets, check_compatibility=check_compatibility)._subset(**kwargs)
 
 
-def zip_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Zip:
+def zip_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Zip:
     """Factory function to create a Zip instance.
 
     Parameters

anemoi/datasets/dates/__init__.py

@@ -10,15 +10,10 @@
 
 import datetime
 import warnings
+from collections.abc import Iterator
 from functools import reduce
 from math import gcd
 from typing import Any
-from typing import Dict
-from typing import Iterator
-from typing import List
-from typing import Optional
-from typing import Tuple
-from typing import Union
 
 # from anemoi.utils.dates import as_datetime
 from anemoi.utils.dates import DateTimes
@@ -29,7 +24,7 @@ from anemoi.utils.hindcasts import HindcastDatesTimes
 from anemoi.utils.humanize import print_dates
 
 
-def extend(x: Union[str, List[Any], Tuple[Any, ...]]) -> Iterator[datetime.datetime]:
+def extend(x: str | list[Any] | tuple[Any, ...]) -> Iterator[datetime.datetime]:
     """Extend a date range or list of dates into individual datetime objects.
 
     Args:
@@ -63,6 +58,8 @@ def extend(x: Union[str, List[Any], Tuple[Any, ...]]) -> Iterator[datetime.datet
 class DatesProvider:
     """Base class for date generation.
 
+    Examples
+    --------
     >>> DatesProvider.from_config(**{"start": "2023-01-01 00:00", "end": "2023-01-02 00:00", "frequency": "1d"}).values
     [datetime.datetime(2023, 1, 1, 0, 0), datetime.datetime(2023, 1, 2, 0, 0)]
 
@@ -86,7 +83,7 @@ class DatesProvider:
     3
     """
 
-    def __init__(self, missing: Optional[List[Union[str, datetime.datetime]]] = None) -> None:
+    def __init__(self, missing: list[str | datetime.datetime] | None = None) -> None:
         """Initialize the DatesProvider with optional missing dates.
 
         Parameters
@@ -168,7 +165,7 @@ class ValuesDates(DatesProvider):
         **kwargs (Any): Additional arguments.
     """
 
-    def __init__(self, values: List[Union[str, datetime.datetime]], **kwargs: Any) -> None:
+    def __init__(self, values: list[str | datetime.datetime], **kwargs: Any) -> None:
         """Initialize ValuesDates with a list of values.
 
         Args:
@@ -188,7 +185,7 @@ class ValuesDates(DatesProvider):
         """
         return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the ValuesDates instance to a dictionary.
 
         Returns
@@ -215,9 +212,9 @@ class StartEndDates(DatesProvider):
 
     def __init__(
         self,
-        start: Union[str, datetime.datetime],
-        end: Union[str, datetime.datetime],
-        frequency: Union[int, str] = 1,
+        start: str | datetime.datetime,
+        end: str | datetime.datetime,
+        frequency: int | str = 1,
         **kwargs: Any,
     ) -> None:
         """Initialize StartEndDates with start, end, and frequency.
@@ -259,7 +256,7 @@ class StartEndDates(DatesProvider):
 
         super().__init__(missing=missing)
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the StartEndDates instance to a dictionary.
 
         Returns
@@ -314,9 +311,9 @@ class HindcastsDates(DatesProvider):
 
     def __init__(
         self,
-        start: Union[str, List[str]],
-        end: Union[str, List[str]],
-        steps: List[int] = [0],
+        start: str | list[str],
+        end: str | list[str],
+        steps: list[int] = [0],
         years: int = 20,
         **kwargs: Any,
     ) -> None:
@@ -403,7 +400,7 @@ class HindcastsDates(DatesProvider):
         """
         return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the HindcastsDates instance to a dictionary.
 
         Returns

anemoi/datasets/dates/groups.py

@@ -12,19 +12,16 @@ import datetime
 import itertools
 from abc import ABC
 from abc import abstractmethod
+from collections.abc import Callable
+from collections.abc import Iterator
 from functools import cached_property
 from typing import Any
-from typing import Callable
-from typing import Iterator
-from typing import List
-from typing import Tuple
-from typing import Union
 
 from anemoi.datasets.dates import DatesProvider
 from anemoi.datasets.dates import as_datetime
 
 
-def _shorten(dates: Union[List[datetime.datetime], Tuple[datetime.datetime, ...]]) -> Union[str, List[str]]:
+def _shorten(dates: list[datetime.datetime] | tuple[datetime.datetime, ...]) -> str | list[str]:
     """Shorten the list of dates for display.
 
     Args:
@@ -43,7 +40,7 @@ def _shorten(dates: Union[List[datetime.datetime], Tuple[datetime.datetime, ...]
 class GroupOfDates:
     """A class to represent a group of dates."""
 
-    def __init__(self, dates: List[datetime.datetime], provider: DatesProvider, partial_ok: bool = False) -> None:
+    def __init__(self, dates: list[datetime.datetime], provider: DatesProvider, partial_ok: bool = False) -> None:
         assert isinstance(provider, DatesProvider), type(provider)
         assert isinstance(dates, list)
 
@@ -197,10 +194,10 @@ class Groups:
 class Filter:
     """A class to filter out missing dates."""
 
-    def __init__(self, missing: List[datetime.datetime]) -> None:
-        self.missing = set(as_datetime(m) for m in missing)
+    def __init__(self, missing: list[datetime.datetime]) -> None:
+        self.missing = {as_datetime(m) for m in missing}
 
-    def __call__(self, dates: List[datetime.datetime]) -> List[datetime.datetime]:
+    def __call__(self, dates: list[datetime.datetime]) -> list[datetime.datetime]:
         """Filter out missing dates from the list of dates.
 
         Args:

anemoi/datasets/dumper.py

@@ -0,0 +1,76 @@
+# (C) Copyright 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.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import datetime
+import io
+import logging
+
+import ruamel.yaml
+
+LOG = logging.getLogger(__name__)
+
+
+def represent_date(dumper, data):
+
+    if isinstance(data, datetime.datetime):
+        if data.tzinfo is None:
+            data = data.replace(tzinfo=datetime.timezone.utc)
+        data = data.astimezone(datetime.timezone.utc)
+        iso_str = data.replace(tzinfo=None).isoformat(timespec="seconds") + "Z"
+    else:
+        iso_str = data.isoformat()
+
+    return dumper.represent_scalar("tag:yaml.org,2002:timestamp", iso_str)
+
+
+# --- Represent multiline strings with | style ---
+def represent_multiline_str(dumper, data):
+    if "\n" in data:
+        return dumper.represent_scalar("tag:yaml.org,2002:str", data.strip(), style="|")
+    return dumper.represent_scalar("tag:yaml.org,2002:str", data)
+
+
+# --- Represent short lists inline (flow style) ---
+def represent_inline_list(dumper, data):
+
+    if not all(isinstance(i, (str, int, float, bool, type(None))) for i in data):
+        return dumper.represent_sequence("tag:yaml.org,2002:seq", data)
+
+    return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True)
+
+
+def yaml_dump(obj, order=None, stream=None, **kwargs):
+
+    if order:
+
+        def _ordering(k):
+            return order.index(k) if k in order else len(order)
+
+        obj = {k: v for k, v in sorted(obj.items(), key=lambda item: _ordering(item[0]))}
+
+    yaml = ruamel.yaml.YAML()
+    yaml.width = 120  # wrap long flow sequences
+
+    yaml.Representer.add_representer(datetime.date, represent_date)
+    yaml.Representer.add_representer(datetime.datetime, represent_date)
+    yaml.Representer.add_representer(str, represent_multiline_str)
+    yaml.Representer.add_representer(list, represent_inline_list)
+
+    data = ruamel.yaml.comments.CommentedMap()
+    for i, (k, v) in enumerate(obj.items()):
+        data[k] = v
+        if i > 0:
+            data.yaml_set_comment_before_after_key(key=k, before="\n")
+
+    if stream:
+        yaml.dump(data, stream=stream, **kwargs)
+
+    stream = io.StringIO()
+    yaml.dump(data, stream=stream, **kwargs)
+    return stream.getvalue()