anemoi-datasets 0.5.15__py3-none-any.whl → 0.5.17__py3-none-any.whl

anemoi/datasets/data/misc.py

@@ -12,21 +12,48 @@ import calendar
 import datetime
 import logging
 from pathlib import PurePath
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
 
 import numpy as np
 import zarr
 from anemoi.utils.config import load_config as load_settings
+from numpy.typing import NDArray
 
-from .dataset import Dataset
+if TYPE_CHECKING:
+    from .dataset import Dataset
 
 LOG = logging.getLogger(__name__)
 
 
-def load_config():
+def load_config() -> Dict[str, Any]:
+    """Load the configuration settings.
+
+    Returns
+    -------
+    Dict[str, Any]
+        The configuration settings.
+    """
     return load_settings(defaults={"datasets": {"named": {}, "path": []}})
 
 
-def add_named_dataset(name, path, **kwargs):
+def add_named_dataset(name: str, path: str, **kwargs: Any) -> None:
+    """Add a named dataset to the configuration.
+
+    Parameters
+    ----------
+    name : str
+        The name of the dataset.
+    path : str
+        The path to the dataset.
+    **kwargs : Any
+        Additional arguments.
+    """
     config = load_config()
     if name["datasets"]["named"]:
         raise ValueError(f"Dataset {name} already exists")
@@ -34,15 +61,37 @@ def add_named_dataset(name, path, **kwargs):
     config["datasets"]["named"][name] = path
 
 
-def add_dataset_path(path):
+def add_dataset_path(path: str) -> None:
+    """Add a dataset path to the configuration.
+
+    Parameters
+    ----------
+    path : str
+        The path to add.
+    """
     config = load_config()
 
     if path not in config["datasets"]["path"]:
         config["datasets"]["path"].append(path)
 
 
-def round_datetime(d, dates, up):
-    """Round up (or down) a datetime to the nearest date in a list of dates"""
+def round_datetime(d: np.datetime64, dates: NDArray[np.datetime64], up: bool) -> np.datetime64:
+    """Round up (or down) a datetime to the nearest date in a list of dates.
+
+    Parameters
+    ----------
+    d : np.datetime64
+        The datetime to round.
+    dates : NDArray[np.datetime64]
+        The list of dates.
+    up : bool
+        Whether to round up.
+
+    Returns
+    -------
+    np.datetime64
+        The rounded datetime.
+    """
     if dates is None or len(dates) == 0:
         return d
 
@@ -58,8 +107,25 @@ def round_datetime(d, dates, up):
     return dates[-1]
 
 
-def _as_date(d, dates, last):
-
+def _as_date(
+    d: Union[int, str, np.datetime64, datetime.date], dates: NDArray[np.datetime64], last: bool
+) -> np.datetime64:
+    """Convert a date to a numpy datetime64 object, rounding to the nearest date in a list of dates.
+
+    Parameters
+    ----------
+    d : Union[int, str, np.datetime64, datetime.date]
+        The date to convert.
+    dates : NDArray[np.datetime64]
+        The list of dates.
+    last : bool
+        Whether to round to the last date.
+
+    Returns
+    -------
+    np.datetime64
+        The converted date.
+    """
     # WARNING,  datetime.datetime is a subclass of datetime.date
     # so we need to check for datetime.datetime first
 
@@ -103,7 +169,7 @@ def _as_date(d, dates, last):
 
     if isinstance(d, str):
 
-        def isfloat(s):
+        def isfloat(s: str) -> bool:
             try:
                 float(s)
                 return True
@@ -164,16 +230,57 @@ def _as_date(d, dates, last):
     raise NotImplementedError(f"Unsupported date: {d} ({type(d)})")
 
 
-def as_first_date(d, dates):
+def as_first_date(d: Union[int, str, np.datetime64, datetime.date], dates: NDArray[np.datetime64]) -> np.datetime64:
+    """Convert a date to the first date in a list of dates.
+
+    Parameters
+    ----------
+    d : Union[int, str, np.datetime64, datetime.date]
+        The date to convert.
+    dates : NDArray[np.datetime64]
+        The list of dates.
+
+    Returns
+    -------
+    np.datetime64
+        The first date.
+    """
     return _as_date(d, dates, last=False)
 
 
-def as_last_date(d, dates):
+def as_last_date(d: Union[int, str, np.datetime64, datetime.date], dates: NDArray[np.datetime64]) -> np.datetime64:
+    """Convert a date to the last date in a list of dates.
+
+    Parameters
+    ----------
+    d : Union[int, str, np.datetime64, datetime.date]
+        The date to convert.
+    dates : NDArray[np.datetime64]
+        The list of dates.
+
+    Returns
+    -------
+    np.datetime64
+        The last date.
+    """
     return _as_date(d, dates, last=True)
 
 
-def _concat_or_join(datasets, kwargs):
+def _concat_or_join(datasets: List["Dataset"], kwargs: Dict[str, Any]) -> Tuple["Dataset", Dict[str, Any]]:
+    """Concatenate or join datasets based on their date ranges.
 
+    Parameters
+    ----------
+    datasets : List[Dataset]
+        The list of datasets.
+    kwargs : Dict[str, Any]
+        Additional arguments.
+
+    Returns
+    -------
+    Tuple[Dataset, Dict[str, Any]]
+        The concatenated or joined dataset and remaining arguments.
+    """
     if "adjust" in kwargs:
         raise ValueError("Cannot use 'adjust' without specifying 'concat' or 'join'")
     datasets, kwargs = _auto_adjust(datasets, kwargs)
@@ -193,7 +300,20 @@ def _concat_or_join(datasets, kwargs):
     return Concat(datasets), kwargs
 
 
-def _open(a):
+def _open(a: Union[str, PurePath, Dict[str, Any], List[Any], Tuple[Any, ...]]) -> "Dataset":
+    """Open a dataset from various input types.
+
+    Parameters
+    ----------
+    a : Union[str, PurePath, Dict[str, Any], List[Any], Tuple[Any, ...]]
+        The input to open.
+
+    Returns
+    -------
+    Dataset
+        The opened dataset.
+    """
+    from .dataset import Dataset
     from .stores import Zarr
     from .stores import zarr_lookup
 
@@ -218,8 +338,27 @@ def _open(a):
     raise NotImplementedError(f"Unsupported argument: {type(a)}")
 
 
-def _auto_adjust(datasets, kwargs, exclude=None):
-
+def _auto_adjust(
+    datasets: List["Dataset"],
+    kwargs: Dict[str, Any],
+    exclude: Optional[List[str]] = None,
+) -> Tuple[List["Dataset"], Dict[str, Any]]:
+    """Automatically adjust datasets based on specified criteria.
+
+    Parameters
+    ----------
+    datasets : List[Dataset]
+        The list of datasets.
+    kwargs : Dict[str, Any]
+        Additional arguments.
+    exclude : Optional[List[str]]
+        List of keys to exclude from adjustment.
+
+    Returns
+    -------
+    Tuple[List[Dataset], Dict[str, Any]]
+        The adjusted datasets and remaining arguments.
+    """
     if "adjust" not in kwargs:
         return datasets, kwargs
 
@@ -296,31 +435,48 @@ def _auto_adjust(datasets, kwargs, exclude=None):
     return datasets, kwargs
 
 
-def _open_dataset(*args, **kwargs):
+def _open_dataset(*args: Any, **kwargs: Any) -> "Dataset":
+    """Open a dataset.
+
+    Parameters
+    ----------
+    *args : Any
+        Positional arguments.
+    **kwargs : Any
+        Keyword arguments.
 
+    Returns
+    -------
+    Dataset
+        The opened dataset.
+    """
     sets = []
     for a in args:
         sets.append(_open(a))
 
     if "xy" in kwargs:
+        # Experimental feature, may be removed
         from .xy import xy_factory
 
         assert not sets, sets
         return xy_factory(args, kwargs).mutate()
 
     if "x" in kwargs and "y" in kwargs:
+        # Experimental feature, may be removed
         from .xy import xy_factory
 
         assert not sets, sets
         return xy_factory(args, kwargs).mutate()
 
     if "zip" in kwargs:
+        # Experimental feature, may be removed
         from .xy import zip_factory
 
         assert not sets, sets
         return zip_factory(args, kwargs).mutate()
 
     if "chain" in kwargs:
+        # Experimental feature, may be removed
         from .unchecked import chain_factory
 
         assert not sets, sets

anemoi/datasets/data/missing.py

@@ -8,14 +8,25 @@
 # nor does it submit to any jurisdiction.
 
 
+import datetime
 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 typing import Union
 
 import numpy as np
+from numpy.typing import NDArray
 
 from anemoi.datasets.create.utils import to_datetime
 from anemoi.datasets.data import MissingDateError
 
+from .dataset import Dataset
+from .dataset import FullIndex
+from .dataset import TupleIndex
 from .debug import Node
 from .debug import debug_indexing
 from .forwards import Forwards
@@ -24,11 +35,30 @@ from .indexing import update_tuple
 
 LOG = logging.getLogger(__name__)
 
+# TODO: Use that class instead of ZarrMissing
 
-class MissingDates(Forwards):
-    # TODO: Use that class instead of ZarrMissing
 
-    def __init__(self, dataset, missing_dates):
+class MissingDates(Forwards):
+    """Handles missing dates in a dataset.
+
+    Attributes
+    ----------
+    dataset : Dataset
+        The dataset object.
+    missing_dates : List[Union[int, str]]
+        List of missing dates.
+    """
+
+    def __init__(self, dataset: Dataset, missing_dates: List[Union[int, str]]) -> None:
+        """Initializes the MissingDates class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset object.
+        missing_dates : List[Union[int, str]]
+            List of missing dates.
+        """
         super().__init__(dataset)
         self.missing_dates = []
 
@@ -56,12 +86,25 @@ class MissingDates(Forwards):
         assert len(self._missing), "No dates to force missing"
 
     @cached_property
-    def missing(self):
+    def missing(self) -> Set[int]:
+        """Returns the set of missing indices."""
         return self._missing.union(self.forward.missing)
 
     @debug_indexing
     @expand_list_indexing
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Retrieves the item at the given index.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index to retrieve.
+
+        Returns
+        -------
+        NDArray[Any]
+            The item at the given index.
+        """
         if isinstance(n, int):
             if n in self.missing:
                 self._report_missing(n)
@@ -94,23 +137,63 @@ class MissingDates(Forwards):
 
         raise TypeError(f"Unsupported index {n} {type(n)}")
 
-    def _report_missing(self, n):
+    def _report_missing(self, n: int) -> None:
+        """Reports a missing date.
+
+        Parameters
+        ----------
+        n : int
+            The index of the missing date.
+        """
         raise MissingDateError(f"Date {self.forward.dates[n]} is missing (index={n})")
 
     @property
-    def reason(self):
+    def reason(self) -> Dict[str, Any]:
+        """Provides the reason for missing dates."""
         return {"missing_dates": self.missing_dates}
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Builds a tree representation of the missing dates.
+
+        Returns
+        -------
+        Node
+            The tree representation of the missing dates.
+        """
         return Node(self, [self.forward.tree()], **self.reason)
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Provides metadata specific to the subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Metadata specific to the subclass.
+        """
         return {"missing_dates": self.missing_dates}
 
 
 class SkipMissingDates(Forwards):
-
-    def __init__(self, dataset, expected_access):
+    """Skips missing dates in a dataset.
+
+    Attributes
+    ----------
+    dataset : Dataset
+        The dataset object.
+    expected_access : Union[int, slice]
+        The expected access pattern.
+    """
+
+    def __init__(self, dataset: Dataset, expected_access: Union[int, slice]) -> None:
+        """Initializes the SkipMissingDates class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset object.
+        expected_access : Union[int, slice]
+            The expected access pattern.
+        """
         super().__init__(dataset)
 
         # if isinstance(expected_access, (tuple, list)):
@@ -141,24 +224,46 @@ class SkipMissingDates(Forwards):
         self.expected_access = expected_access
         self.indices = indices
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Returns the length of the indices.
+
+        Returns
+        -------
+        int
+            The length of the indices.
+        """
         return len(self.indices)
 
     @property
-    def start_date(self):
+    def start_date(self) -> np.datetime64:
+        """Returns the start date."""
         return self.forward.start_date
 
     @property
-    def end_date(self):
+    def end_date(self) -> np.datetime64:
+        """Returns the end date."""
         return self.forward.end_date
 
     @property
-    def dates(self):
+    def dates(self) -> NDArray[np.datetime64]:
+        """Not implemented. Raises an error."""
         raise NotImplementedError("SkipMissingDates.dates")
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Retrieves a tuple of items at the given index.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The index to retrieve.
+
+        Returns
+        -------
+        NDArray[Any]
+            The tuple of items at the given index.
+        """
 
         def _get_one(n):
             result = []
@@ -180,13 +285,37 @@ class SkipMissingDates(Forwards):
         return tuple(np.stack(_) for _ in result)
 
     @debug_indexing
-    def _get_slice(self, s):
+    def _get_slice(self, s: slice) -> Tuple[NDArray[Any], ...]:
+        """Retrieves a slice of items.
+
+        Parameters
+        ----------
+        s : slice
+            The slice to retrieve.
+
+        Returns
+        -------
+        Tuple[NDArray[Any], ...]
+            The slice of items.
+        """
         values = [self[i] for i in range(*s.indices(self._len))]
         result = [_ for _ in zip(*values)]
         return tuple(np.stack(_) for _ in result)
 
     @debug_indexing
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> Tuple[NDArray[Any], ...]:
+        """Retrieves the item at the given index.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index to retrieve.
+
+        Returns
+        -------
+        Tuple[NDArray[Any], ...]
+            The item at the given index.
+        """
         if isinstance(n, tuple):
             return self._get_tuple(n)
 
@@ -196,19 +325,56 @@ class SkipMissingDates(Forwards):
         return tuple(self.forward[i] for i in self.indices[n])
 
     @property
-    def frequency(self):
+    def frequency(self) -> datetime.timedelta:
+        """Returns the frequency of the dataset."""
         return self.forward.frequency
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Builds a tree representation of the skipped missing dates.
+
+        Returns
+        -------
+        Node
+            The tree representation of the skipped missing dates.
+        """
         return Node(self, [self.forward.tree()], expected_access=self.expected_access)
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Provides metadata specific to the subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Metadata specific to the subclass.
+        """
         return {"expected_access": self.expected_access}
 
 
 class MissingDataset(Forwards):
-
-    def __init__(self, dataset, start, end):
+    """Represents a dataset with missing dates.
+
+    Attributes
+    ----------
+    dataset : Dataset
+        The dataset object.
+    start : np.datetime64
+        The start date.
+    end : np.datetime64
+        The end date.
+    """
+
+    def __init__(self, dataset: Dataset, start: np.datetime64, end: np.datetime64) -> None:
+        """Initializes the MissingDataset class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset object.
+        start : np.datetime64
+            The start date.
+        end : np.datetime64
+            The end date.
+        """
         super().__init__(dataset)
         self.start = start
         self.end = end
@@ -222,22 +388,60 @@ class MissingDataset(Forwards):
         self._dates = np.array(dates, dtype="datetime64")
         self._missing = set(range(len(dates)))
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Returns the length of the dates.
+
+        Returns
+        -------
+        int
+            The length of the dates.
+        """
         return len(self._dates)
 
     @property
-    def dates(self):
+    def dates(self) -> NDArray[np.datetime64]:
+        """Returns the dates of the dataset."""
         return self._dates
 
     @property
-    def missing(self):
+    def missing(self) -> Set[int]:
+        """Returns the set of missing indices."""
         return self._missing
 
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Raises an error for missing dates.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index to retrieve.
+
+        Raises
+        ------
+        MissingDateError
+            If the date is missing.
+
+        Returns:
+            NDArray[Any]: The data at the specified index.
+        """
         raise MissingDateError(f"Date {self.dates[n]} is missing (index={n})")
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Builds a tree representation of the missing dataset.
+
+        Returns
+        -------
+        Node
+            The tree representation of the missing dataset.
+        """
         return Node(self, [self.forward.tree()], start=self.start, end=self.end)
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Provides metadata specific to the subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Metadata specific to the subclass.
+        """
         return {"start": self.start, "end": self.end}