anemoi-datasets 0.5.16__py3-none-any.whl → 0.5.18__py3-none-any.whl

anemoi/datasets/create/check.py

@@ -8,25 +8,48 @@
 # nor does it submit to any jurisdiction.
 
 
+import datetime
 import logging
 import re
 import warnings
+from typing import Any
+from typing import Callable
+from typing import Optional
+from typing import Union
 
 import numpy as np
 from anemoi.utils.dates import frequency_to_string
+from numpy.typing import NDArray
 
 LOG = logging.getLogger(__name__)
 
 
 class DatasetName:
+    """Class to validate and parse dataset names according to naming conventions."""
+
     def __init__(
         self,
-        name,
-        resolution=None,
-        start_date=None,
-        end_date=None,
-        frequency=None,
+        name: str,
+        resolution: Optional[str] = None,
+        start_date: Optional[datetime.date] = None,
+        end_date: Optional[datetime.date] = None,
+        frequency: Optional[datetime.timedelta] = None,
     ):
+        """Initialize a DatasetName instance.
+
+        Parameters
+        ----------
+        name : str
+            The name of the dataset.
+        resolution : Optional[str], optional
+            The resolution of the dataset.
+        start_date : Optional[datetime.date], optional
+            The start date of the dataset.
+        end_date : Optional[datetime.date], optional
+            The end date of the dataset.
+        frequency : Optional[datetime.timedelta], optional
+            The frequency of the dataset.
+        """
         self.name = name
         self.parsed = self._parse(name)
         print("---------------")
@@ -45,19 +68,39 @@ class DatasetName:
             self.messages.append(f"{self} is parsed as :" + "/".join(f"{k}={v}" for k, v in self.parsed.items()))
 
     @property
-    def error_message(self):
+    def error_message(self) -> str:
+        """Generate an error message based on the collected messages."""
         out = " And ".join(self.messages)
         if out:
-            out = out[0].upper() + out[1:]
+            out[0].upper() + out[1:]
         return out
 
-    def raise_if_not_valid(self, print=print):
+    def raise_if_not_valid(self, print: Callable = print) -> None:
+        """Raise a ValueError if the dataset name is not valid.
+
+        Parameters
+        ----------
+        print : Callable
+            The function to use for printing messages.
+        """
         if self.messages:
             for m in self.messages:
                 print(m)
             raise ValueError(self.error_message)
 
-    def _parse(self, name):
+    def _parse(self, name: str) -> dict:
+        """Parse the dataset name into its components.
+
+        Parameters
+        ----------
+        name : str
+            The name of the dataset.
+
+        Returns
+        -------
+        dict
+            The parsed components of the dataset name.
+        """
         pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h|\d+m)-v(\d+)-?([a-zA-Z0-9-]+)?$"
         match = re.match(pattern, name)
 
@@ -81,10 +124,12 @@ class DatasetName:
 
         return parsed
 
-    def __str__(self):
+    def __str__(self) -> str:
+        """Return the string representation of the dataset name."""
         return self.name
 
-    def check_parsed(self):
+    def check_parsed(self) -> None:
+        """Check if the dataset name was parsed correctly."""
         if not self.parsed:
             self.messages.append(
                 f"the dataset name {self} does not follow naming convention. "
@@ -92,7 +137,14 @@ class DatasetName:
                 "https://anemoi-registry.readthedocs.io/en/latest/naming-conventions.html"
             )
 
-    def check_resolution(self, resolution):
+    def check_resolution(self, resolution: Optional[str]) -> None:
+        """Check if the resolution matches the expected format.
+
+        Parameters
+        ----------
+        resolution : str or None
+            The expected resolution.
+        """
         if self.parsed.get("resolution") and self.parsed["resolution"][0] not in "0123456789on":
             self.messages.append(
                 f"the resolution {self.parsed['resolution'] } should start "
@@ -105,42 +157,97 @@ class DatasetName:
         self._check_missing("resolution", resolution_str)
         self._check_mismatch("resolution", resolution_str)
 
-    def check_frequency(self, frequency):
+    def check_frequency(self, frequency: Optional[datetime.timedelta]) -> None:
+        """Check if the frequency matches the expected format.
+
+        Parameters
+        ----------
+        frequency : datetime.timedelta or None
+            The expected frequency.
+        """
         if frequency is None:
             return
         frequency_str = frequency_to_string(frequency)
         self._check_missing("frequency", frequency_str)
         self._check_mismatch("frequency", frequency_str)
 
-    def check_start_date(self, start_date):
+    def check_start_date(self, start_date: Optional[datetime.date]) -> None:
+        """Check if the start date matches the expected format.
+
+        Parameters
+        ----------
+        start_date : datetime.date or None
+            The expected start date.
+        """
         if start_date is None:
             return
         start_date_str = str(start_date.year)
         self._check_missing("start_date", start_date_str)
         self._check_mismatch("start_date", start_date_str)
 
-    def check_end_date(self, end_date):
+    def check_end_date(self, end_date: Optional[datetime.date]) -> None:
+        """Check if the end date matches the expected format.
+
+        Parameters
+        ----------
+        end_date : datetime.date or None
+            The expected end date.
+        """
         if end_date is None:
             return
         end_date_str = str(end_date.year)
         self._check_missing("end_date", end_date_str)
         self._check_mismatch("end_date", end_date_str)
 
-    def _check_missing(self, key, value):
+    def _check_missing(self, key: str, value: str) -> None:
+        """Check if a component is missing from the dataset name.
+
+        Parameters
+        ----------
+        key : str
+            The component key.
+        value : str
+            The expected value.
+        """
         if value not in self.name:
             self.messages.append(f"the {key} is {value}, but is missing in {self.name}.")
 
-    def _check_mismatch(self, key, value):
+    def _check_mismatch(self, key: str, value: str) -> None:
+        """Check if a component value mismatches the expected value.
+
+        Parameters
+        ----------
+        key : str
+            The component key.
+        value : str
+            The expected value.
+        """
         if self.parsed.get(key) and self.parsed[key] != value:
             self.messages.append(f"the {key} is {value}, but is {self.parsed[key]} in {self.name}.")
 
 
 class StatisticsValueError(ValueError):
-    pass
+    """Custom error for statistics value issues."""
 
+    pass
 
-def check_data_values(arr, *, name: str, log=[], allow_nans=False):
 
+def check_data_values(
+    arr: NDArray[Any], *, name: str, log: list = [], allow_nans: Union[bool, list, set, tuple, dict] = False
+) -> None:
+    """Check the values in the data array for validity.
+
+    Parameters
+    ----------
+    arr : NDArray[Any]
+        The data array to check.
+    name : str
+        The name of the data array.
+    log : list, optional
+        A list to log messages.
+    allow_nans : bool or list or set or tuple or dict, optional
+        Whether to allow NaNs in the data array.
+    """
     shape = arr.shape
 
     if (isinstance(allow_nans, (set, list, tuple, dict)) and name in allow_nans) or allow_nans:
@@ -182,7 +289,22 @@ def check_data_values(arr, *, name: str, log=[], allow_nans=False):
             )
 
 
-def check_stats(minimum, maximum, mean, msg, **kwargs):
+def check_stats(minimum: float, maximum: float, mean: float, msg: str, **kwargs: Any) -> None:
+    """Check if the mean value is within the min/max interval.
+
+    Parameters
+    ----------
+    minimum : float
+        The minimum value.
+    maximum : float
+        The maximum value.
+    mean : float
+        The mean value.
+    msg : str
+        The message to include in the error.
+    **kwargs : Any
+        Additional keyword arguments.
+    """
     tolerance = (abs(minimum) + abs(maximum)) * 0.01
     if (mean - minimum < -tolerance) or (mean - minimum < -tolerance):
         raise StatisticsValueError(

anemoi/datasets/create/chunks.py

@@ -9,6 +9,7 @@
 
 import logging
 import warnings
+from typing import Union
 
 LOG = logging.getLogger(__name__)
 
@@ -16,7 +17,35 @@ ALL = object()
 
 
 class ChunkFilter:
-    def __init__(self, *, parts, total):
+    """A filter to determine which chunks to process based on the specified parts.
+
+    Attributes
+    ----------
+    total : int
+        The total number of chunks.
+    allowed : object or list
+        The chunks that are allowed to be processed.
+    """
+
+    def __init__(self, *, parts: Union[str, list], total: int):
+        """Initializes the ChunkFilter with the given parts and total number of chunks.
+
+        Parameters
+        ----------
+        parts : str or list
+            The parts to process, specified as 'i/n' or a list of such strings.
+        total : int
+            The total number of chunks.
+
+        Raises
+        ------
+        ValueError
+            If the parts format is invalid.
+        AssertionError
+            If the chunk number is invalid.
+        Warning
+            If the number of chunks is larger than the total number of chunks.
+        """
         self.total = total
 
         if isinstance(parts, list):
@@ -62,7 +91,24 @@ class ChunkFilter:
 
         self.allowed = parts
 
-    def __call__(self, i):
+    def __call__(self, i: int) -> bool:
+        """Checks if the given chunk number is allowed to be processed.
+
+        Parameters
+        ----------
+        i : int
+            The chunk number to check.
+
+        Returns
+        -------
+        bool
+            True if the chunk is allowed, False otherwise.
+
+        Raises
+        ------
+        AssertionError
+            If the chunk number is invalid.
+        """
         if i < 0 or i >= self.total:
             raise AssertionError(f"Invalid chunk number {i}. Must be between 0 and {self.total - 1}.")
 
@@ -70,10 +116,24 @@ class ChunkFilter:
             return True
         return i in self.allowed
 
-    def __iter__(self):
+    def __iter__(self) -> iter:
+        """Iterates over the allowed chunks.
+
+        Yields
+        ------
+        int
+            The next allowed chunk number.
+        """
         for i in range(self.total):
             if self(i):
                 yield i
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Returns the number of allowed chunks.
+
+        Returns
+        -------
+        int
+            The number of allowed chunks.
+        """
         return len([_ for _ in self])

anemoi/datasets/create/config.py

@@ -11,6 +11,9 @@ import datetime
 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
@@ -22,13 +25,41 @@ from anemoi.datasets.dates.groups import Groups
 LOG = logging.getLogger(__name__)
 
 
-def _get_first_key_if_dict(x):
+def _get_first_key_if_dict(x: Union[str, dict]) -> str:
+    """Returns the first key if the input is a dictionary, otherwise returns the input string.
+
+    Parameters
+    ----------
+    x : str or dict
+        Input string or dictionary.
+
+    Returns
+    -------
+    str
+        The first key if input is a dictionary, otherwise the input string.
+    """
     if isinstance(x, str):
         return x
     return list(x.keys())[0]
 
 
-def ensure_element_in_list(lst, elt, index):
+def ensure_element_in_list(lst: list, elt: str, index: int) -> list:
+    """Ensures that a specified element is present at a given index in a list.
+
+    Parameters
+    ----------
+    lst : list
+        The list to check.
+    elt : str
+        The element to ensure is in the list.
+    index : int
+        The index at which the element should be present.
+
+    Returns
+    -------
+    list
+        The modified list with the element at the specified index.
+    """
     if elt in lst:
         assert lst[index] == elt
         return lst
@@ -41,7 +72,23 @@ def ensure_element_in_list(lst, elt, index):
     return lst[:index] + [elt] + lst[index:]
 
 
-def check_dict_value_and_set(dic, key, value):
+def check_dict_value_and_set(dic: dict, key: str, value: Any) -> None:
+    """Checks if a dictionary contains a specific key-value pair and sets it if not present.
+
+    Parameters
+    ----------
+    dic : dict
+        The dictionary to check.
+    key : str
+        The key to check in the dictionary.
+    value : Any
+        The value to set if the key is not present.
+
+    Raises
+    ------
+    ValueError
+        If the key is present but with a different value.
+    """
     if key in dic:
         if dic[key] == value:
             return
@@ -50,7 +97,19 @@ def check_dict_value_and_set(dic, key, value):
     dic[key] = value
 
 
-def resolve_includes(config):
+def resolve_includes(config: Union[dict, list]) -> Union[dict, list]:
+    """Resolves '<<' includes in a configuration dictionary or list.
+
+    Parameters
+    ----------
+    config : dict or list
+        The configuration to resolve includes for.
+
+    Returns
+    -------
+    dict or list
+        The configuration with includes resolved.
+    """
     if isinstance(config, list):
         return [resolve_includes(c) for c in config]
     if isinstance(config, dict):
@@ -62,7 +121,18 @@ def resolve_includes(config):
 
 
 class Config(DotDict):
-    def __init__(self, config=None, **kwargs):
+    """Configuration class that extends DotDict to handle configuration loading and processing."""
+
+    def __init__(self, config: Optional[Union[str, dict]] = None, **kwargs):
+        """Initializes the Config object.
+
+        Parameters
+        ----------
+        config : str or dict, optional
+            Path to the configuration file or a dictionary. Defaults to None.
+        **kwargs
+            Additional keyword arguments to update the configuration.
+        """
         if isinstance(config, str):
             self.config_path = os.path.realpath(config)
             config = load_any_dict_format(config)
@@ -74,7 +144,18 @@ class Config(DotDict):
 
 
 class OutputSpecs:
-    def __init__(self, config, parent):
+    """Class to handle output specifications for datasets."""
+
+    def __init__(self, config: Config, parent: Any):
+        """Initializes the OutputSpecs object.
+
+        Parameters
+        ----------
+        config : Config
+            The configuration object.
+        parent : Any
+            The parent object.
+        """
         self.config = config
         if "order_by" in config:
             assert isinstance(config.order_by, dict), config.order_by
@@ -82,15 +163,28 @@ class OutputSpecs:
         self.parent = parent
 
     @property
-    def dtype(self):
+    def dtype(self) -> str:
+        """Returns the data type for the output."""
         return self.config.dtype
 
     @property
-    def order_by_as_list(self):
-        # this is used when an ordered dict is not supported (e.g. zarr attributes)
+    def order_by_as_list(self) -> list[dict]:
+        """Returns the order_by configuration as a list of dictionaries."""
         return [{k: v} for k, v in self.config.order_by.items()]
 
-    def get_chunking(self, coords):
+    def get_chunking(self, coords: dict) -> tuple:
+        """Returns the chunking configuration based on coordinates.
+
+        Parameters
+        ----------
+        coords : dict
+            The coordinates dictionary.
+
+        Returns
+        -------
+        tuple
+            The chunking configuration.
+        """
         user = deepcopy(self.config.chunking)
         chunks = []
         for k, v in coords.items():
@@ -105,25 +199,41 @@ class OutputSpecs:
         return tuple(chunks)
 
     @property
-    def order_by(self):
+    def order_by(self) -> dict:
+        """Returns the order_by configuration."""
         return self.config.order_by
 
     @property
-    def remapping(self):
+    def remapping(self) -> dict:
+        """Returns the remapping configuration."""
         return self.config.remapping
 
     @property
-    def flatten_grid(self):
+    def flatten_grid(self) -> bool:
+        """Returns whether the grid should be flattened."""
         return self.config.flatten_grid
 
     @property
-    def statistics(self):
+    def statistics(self) -> str:
+        """Returns the statistics configuration."""
         return self.config.statistics
 
 
 class LoadersConfig(Config):
-    def __init__(self, config, *args, **kwargs):
-
+    """Configuration class for dataset loaders."""
+
+    def __init__(self, config: dict, *args, **kwargs):
+        """Initializes the LoadersConfig object.
+
+        Parameters
+        ----------
+        config : dict
+            The configuration dictionary.
+        *args
+            Additional positional arguments.
+        **kwargs
+            Additional keyword arguments.
+        """
         super().__init__(config, *args, **kwargs)
 
         # TODO: should use a json schema to validate the config
@@ -178,11 +288,30 @@ class LoadersConfig(Config):
 
         self.reading_chunks = self.get("reading_chunks")
 
-    def get_serialisable_dict(self):
+    def get_serialisable_dict(self) -> dict:
+        """Returns a serializable dictionary representation of the configuration.
+
+        Returns
+        -------
+        dict
+            The serializable dictionary.
+        """
         return _prepare_serialisation(self)
 
 
-def _prepare_serialisation(o):
+def _prepare_serialisation(o: Any) -> Any:
+    """Prepares an object for serialization.
+
+    Parameters
+    ----------
+    o : Any
+        The object to prepare.
+
+    Returns
+    -------
+    Any
+        The prepared object.
+    """
     if isinstance(o, dict):
         dic = {}
         for k, v in o.items():
@@ -212,7 +341,14 @@ def _prepare_serialisation(o):
     return str(o)
 
 
-def set_to_test_mode(cfg):
+def set_to_test_mode(cfg: dict) -> None:
+    """Modifies the configuration to run in test mode.
+
+    Parameters
+    ----------
+    cfg : dict
+        The configuration dictionary.
+    """
     NUMBER_OF_DATES = 4
 
     LOG.warning(f"Running in test mode. Changing the list of dates to use only {NUMBER_OF_DATES}.")
@@ -251,7 +387,21 @@ def set_to_test_mode(cfg):
     set_element_to_test(cfg)
 
 
-def loader_config(config, is_test=False):
+def loader_config(config: dict, is_test: bool = False) -> LoadersConfig:
+    """Loads and validates the configuration for dataset loaders.
+
+    Parameters
+    ----------
+    config : dict
+        The configuration dictionary.
+    is_test : bool, optional
+        Whether to run in test mode. Defaults to False.
+
+    Returns
+    -------
+    LoadersConfig
+        The validated configuration object.
+    """
     config = Config(config)
     if is_test:
         set_to_test_mode(config)
@@ -273,5 +423,19 @@ def loader_config(config, is_test=False):
     return copy
 
 
-def build_output(*args, **kwargs):
+def build_output(*args, **kwargs) -> OutputSpecs:
+    """Builds the output specifications.
+
+    Parameters
+    ----------
+    *args
+        Additional positional arguments.
+    **kwargs
+        Additional keyword arguments.
+
+    Returns
+    -------
+    OutputSpecs
+        The output specifications object.
+    """
     return OutputSpecs(*args, **kwargs)