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

anemoi/datasets/create/input/{result.py → result/field.py}

@@ -15,10 +15,6 @@ from collections import defaultdict
 from functools import cached_property
 from typing import Any
 from typing import DefaultDict
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Tuple
 
 import numpy as np
 from anemoi.utils.dates import as_timedelta
@@ -26,14 +22,12 @@ from anemoi.utils.humanize import seconds_to_human
 from anemoi.utils.humanize import shorten_list
 from earthkit.data.core.order import build_remapping
 
-from .action import ActionContext
-from .trace import trace
-from .trace import trace_datasource
+from . import Result
 
 LOG = logging.getLogger(__name__)
 
 
-def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
+def _fields_metatata(variables: tuple[str, ...], cube: Any) -> dict[str, Any]:
     """Retrieve metadata for the given variables and cube.
 
     Parameters
@@ -50,7 +44,7 @@ def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
     """
     assert isinstance(variables, tuple), variables
 
-    KNOWN: Dict[str, Dict[str, bool]] = {
+    KNOWN: dict[str, dict[str, bool]] = {
         "cos_julian_day": dict(computed_forcing=True, constant_in_time=False),
         "cos_latitude": dict(computed_forcing=True, constant_in_time=True),
         "cos_local_time": dict(computed_forcing=True, constant_in_time=False),
@@ -65,9 +59,9 @@ def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
         "sin_longitude": dict(computed_forcing=True, constant_in_time=True),
     }
 
-    def _merge(md1: Dict[str, Any], md2: Dict[str, Any]) -> Dict[str, Any]:
+    def _merge(md1: dict[str, Any], md2: dict[str, Any]) -> dict[str, Any]:
         assert set(md1.keys()) == set(md2.keys()), (set(md1.keys()), set(md2.keys()))
-        result: Dict[str, Any] = {}
+        result: dict[str, Any] = {}
         for k in md1.keys():
             v1 = md1[k]
             v2 = md2[k]
@@ -90,10 +84,10 @@ def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
 
         return result
 
-    mars: Dict[str, Any] = {}
-    other: DefaultDict[str, Dict[str, Any]] = defaultdict(dict)
+    mars: dict[str, Any] = {}
+    other: DefaultDict[str, dict[str, Any]] = defaultdict(dict)
     i: int = -1
-    date: Optional[str] = None
+    date: str | None = None
     for c in cube.iterate_cubelets():
 
         if date is None:
@@ -137,7 +131,7 @@ def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
 
         if startStep != endStep:
             # https://codes.ecmwf.int/grib/format/grib2/ctables/4/10/
-            TYPE_OF_STATISTICAL_PROCESSING: Dict[Optional[int], Optional[str]] = {
+            TYPE_OF_STATISTICAL_PROCESSING: dict[int | None, str | None] = {
                 None: None,
                 0: "average",
                 1: "accumulation",
@@ -157,12 +151,12 @@ def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
 
             # https://codes.ecmwf.int/grib/format/grib1/ctable/5/
 
-            TIME_RANGE_INDICATOR: Dict[int, str] = {
+            TIME_RANGE_INDICATOR: dict[int, str] = {
                 4: "accumulation",
                 3: "average",
             }
 
-            STEP_TYPE_FOR_CONVERSION: Dict[str, str] = {
+            STEP_TYPE_FOR_CONVERSION: dict[str, str] = {
                 "min": "minimum",
                 "max": "maximum",
                 "accum": "accumulation",
@@ -172,7 +166,7 @@ def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
             # A few patches
             #
 
-            PATCHES: Dict[str, str] = {
+            PATCHES: dict[str, str] = {
                 "10fg6": "maximum",
                 "mntpr3": "minimum",  # Not in param db
                 "mntpr6": "minimum",  # Not in param db
@@ -211,7 +205,7 @@ def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
         else:
             mars[variables[i]] = md
 
-    result: Dict[str, Dict[str, Any]] = {}
+    result: dict[str, dict[str, Any]] = {}
     for k, v in mars.items():
         result[k] = dict(mars=v) if v else {}
         result[k].update(other[k])
@@ -222,7 +216,7 @@ def _fields_metatata(variables: Tuple[str, ...], cube: Any) -> Dict[str, Any]:
     return result
 
 
-def _data_request(data: Any) -> Dict[str, Any]:
+def _data_request(data: Any) -> dict[str, Any]:
     """Build a data request dictionary from the given data.
 
     Parameters
@@ -235,12 +229,12 @@ def _data_request(data: Any) -> Dict[str, Any]:
     dict
         The data request dictionary.
     """
-    date: Optional[Any] = None
+    date: Any | None = None
     params_levels: DefaultDict[str, set] = defaultdict(set)
     params_steps: DefaultDict[str, set] = defaultdict(set)
 
-    area: Optional[Any] = None
-    grid: Optional[Any] = None
+    area: Any | None = None
+    grid: Any | None = None
 
     for field in data:
         try:
@@ -270,8 +264,8 @@ def _data_request(data: Any) -> Dict[str, Any]:
         except Exception:
             LOG.error(f"Error in retrieving metadata (cannot build data request info) for {field}", exc_info=True)
 
-    def sort(old_dic: DefaultDict[str, set]) -> Dict[str, List[Any]]:
-        new_dic: Dict[str, List[Any]] = {}
+    def sort(old_dic: DefaultDict[str, set]) -> dict[str, list[Any]]:
+        new_dic: dict[str, list[Any]] = {}
         for k, v in old_dic.items():
             new_dic[k] = sorted(list(v))
         return new_dic
@@ -282,43 +276,25 @@ def _data_request(data: Any) -> Dict[str, Any]:
     return dict(param_level=params_levels, param_step=params_steps, area=area, grid=grid)
 
 
-class Result:
+class FieldResult(Result):
     """Class to represent the result of an action in the dataset creation process."""
 
     empty: bool = False
     _coords_already_built: bool = False
 
-    def __init__(self, context: ActionContext, action_path: List[str], dates: Any) -> None:
-        """Initialize a Result instance.
+    def __init__(self, context: Any, datasource: Any) -> None:
 
-        Parameters
-        ----------
-        context : ActionContext
-            The context in which the result exists.
-        action_path : list of str
-            The action path.
-        dates : Any
-            The dates associated with the result.
-        """
         from anemoi.datasets.dates.groups import GroupOfDates
 
-        assert isinstance(dates, GroupOfDates), dates
-
-        assert isinstance(context, ActionContext), type(context)
-        assert isinstance(action_path, list), action_path
-
         self.context: Any = context
-        self.group_of_dates: Any = dates
-        self.action_path: List[str] = action_path
-
-    @property
-    @trace_datasource
-    def datasource(self) -> Any:
-        """Retrieve the data source for the result."""
-        self._raise_not_implemented()
+        self.datasource = datasource
+        self.group_of_dates = context.argument
+        assert isinstance(
+            self.group_of_dates, GroupOfDates
+        ), f"Expected group_of_dates to be a GroupOfDates, got {type(self.group_of_dates)}: {self.group_of_dates}"
 
     @property
-    def data_request(self) -> Dict[str, Any]:
+    def data_request(self) -> dict[str, Any]:
         """Returns a dictionary with the parameters needed to retrieve the data."""
         return _data_request(self.datasource)
 
@@ -330,7 +306,7 @@ class Result:
         Any
             The data cube.
         """
-        trace("🧊", f"getting cube from {self.__class__.__name__}")
+
         ds: Any = self.datasource
 
         remapping: Any = self.context.remapping
@@ -340,7 +316,7 @@ class Result:
         LOG.debug("Sorting dataset %s %s", dict(order_by), remapping)
         assert order_by, order_by
 
-        patches: Dict[str, Dict[Optional[Any], int]] = {"number": {None: 0}}
+        patches: dict[str, dict[Any | None, int]] = {"number": {None: 0}}
 
         try:
             cube: Any = ds.cube(
@@ -377,7 +353,7 @@ class Result:
         patches : Any
             The patches configuration.
         """
-        METADATA: Tuple[str, ...] = (
+        METADATA: tuple[str, ...] = (
             "date",
             "time",
             "step",
@@ -402,7 +378,7 @@ class Result:
         # print("Executing", self.action_path)
         # print("Dates:", compress_dates(self.dates))
 
-        names: List[str] = []
+        names: list[str] = []
         for a in args:
             if isinstance(a, str):
                 names.append(a)
@@ -421,7 +397,7 @@ class Result:
                 print("     ", n)
 
         print()
-        user_shape: Tuple[int, ...] = tuple(len(v) for k, v in user_coords.items())
+        user_shape: tuple[int, ...] = tuple(len(v) for k, v in user_coords.items())
         print("Shape of the hypercube           :", user_shape)
         print(
             "Number of expected fields        :", math.prod(user_shape), "=", " x ".join([str(i) for i in user_shape])
@@ -523,66 +499,6 @@ class Result:
         print()
         exit(1)
 
-    def _repr(self, *args: Any, _indent_: str = "\n", **kwargs: Any) -> str:
-        """Return the string representation of the Result instance.
-
-        Parameters
-        ----------
-        args : Any
-            Additional positional arguments.
-        _indent_ : str
-            Indentation string.
-        kwargs : Any
-            Additional keyword arguments.
-
-        Returns
-        -------
-        str
-            The string representation.
-        """
-        more: str = ",".join([str(a)[:5000] for a in args])
-        more += ",".join([f"{k}={v}"[:5000] for k, v in kwargs.items()])
-
-        dates: str = " no-dates"
-        if self.group_of_dates is not None:
-            dates = f" {len(self.group_of_dates)} dates"
-            dates += " ("
-            dates += "/".join(d.strftime("%Y-%m-%dT%H:%M") for d in self.group_of_dates)
-            if len(dates) > 100:
-                dates = dates[:100] + "..."
-            dates += ")"
-
-        more = more[:5000]
-        txt: str = f"{self.__class__.__name__}:{dates}{_indent_}{more}"
-        if _indent_:
-            txt = txt.replace("\n", "\n  ")
-        return txt
-
-    def __repr__(self) -> str:
-        """Return the string representation of the Result instance."""
-        return self._repr()
-
-    def _raise_not_implemented(self) -> None:
-        """Raise a NotImplementedError indicating the method is not implemented."""
-        raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
-
-    def _trace_datasource(self, *args: Any, **kwargs: Any) -> str:
-        """Trace the data source for the result.
-
-        Parameters
-        ----------
-        args : Any
-            Additional positional arguments.
-        kwargs : Any
-            Additional keyword arguments.
-
-        Returns
-        -------
-        str
-            The trace string.
-        """
-        return f"{self.__class__.__name__}({self.group_of_dates})"
-
     def build_coords(self) -> None:
         """Build the coordinates for the result."""
         if self._coords_already_built:
@@ -643,13 +559,13 @@ class Result:
         self._coords_already_built: bool = True
 
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Retrieve the variables for the result."""
         self.build_coords()
         return self._variables
 
     @property
-    def variables_metadata(self) -> Dict[str, Any]:
+    def variables_metadata(self) -> dict[str, Any]:
         """Retrieve the metadata for the variables."""
         return _fields_metatata(self.variables, self._cube)
 
@@ -690,7 +606,7 @@ class Result:
         return self._proj_string
 
     @cached_property
-    def shape(self) -> List[int]:
+    def shape(self) -> list[int]:
         """Retrieve the shape of the result."""
         return [
             len(self.group_of_dates),
@@ -700,7 +616,7 @@ class Result:
         ]
 
     @cached_property
-    def coords(self) -> Dict[str, Any]:
+    def coords(self) -> dict[str, Any]:
         """Retrieve the coordinates of the result."""
         return {
             "dates": list(self.group_of_dates),

anemoi/datasets/create/input/trace.py

@@ -11,9 +11,9 @@
 import logging
 import textwrap
 import threading
+from collections.abc import Callable
 from functools import wraps
 from typing import Any
-from typing import Callable
 
 LOG = logging.getLogger(__name__)
 

anemoi/datasets/create/patch.py

@@ -10,14 +10,13 @@
 import json
 import logging
 import os
-from typing import Union
 
 import zarr
 
 LOG = logging.getLogger(__name__)
 
 
-def fix_order_by(order_by: Union[dict, list]) -> list[dict]:
+def fix_order_by(order_by: dict | list) -> list[dict]:
     """Fix the order_by attribute to ensure it is a list of dictionaries.
 
     Parameters

anemoi/datasets/create/persistent.py

@@ -16,9 +16,8 @@ import os
 import pickle
 import shutil
 import socket
+from collections.abc import Iterator
 from typing import Any
-from typing import Iterator
-from typing import Tuple
 
 import numpy as np
 from anemoi.utils.provenance import gather_provenance_info
@@ -198,7 +197,7 @@ class BufferedPersistentDict(PersistentDict):
         self.elements = []
         self.keys = []
 
-    def items(self) -> Iterator[Tuple[Any, Any]]:
+    def items(self) -> Iterator[tuple[Any, Any]]:
         """Yield items stored in the BufferedPersistentDict.
 
         Yields
@@ -207,8 +206,7 @@ class BufferedPersistentDict(PersistentDict):
             An iterator over the items.
         """
         for keys, elements in self.storage.items():
-            for key, elt in zip(keys, elements):
-                yield key, elt
+            yield from zip(keys, elements)
 
     def delete(self) -> None:
         """Delete the storage directory and its contents."""

anemoi/datasets/create/size.py

@@ -10,8 +10,6 @@
 
 import logging
 import os
-from typing import Dict
-from typing import Optional
 
 import tqdm
 from anemoi.utils.humanize import bytes_to_human
@@ -19,7 +17,7 @@ from anemoi.utils.humanize import bytes_to_human
 LOG = logging.getLogger(__name__)
 
 
-def compute_directory_sizes(path: str) -> Optional[Dict[str, int]]:
+def compute_directory_sizes(path: str) -> dict[str, int] | None:
     """Computes the total size and number of files in a directory.
 
     Parameters