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

anemoi/datasets/data/xy.py

@@ -10,7 +10,14 @@
 
 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
 from .debug import Node
 from .forwards import Combined
 from .misc import _auto_adjust
@@ -20,94 +27,217 @@ LOG = logging.getLogger(__name__)
 
 
 class ZipBase(Combined):
-
-    def __init__(self, datasets, check_compatibility=True):
+    """Base class for handling zipped datasets."""
+
+    def __init__(self, datasets: List[Any], check_compatibility: bool = True) -> None:
+        """Initialize ZipBase with a list of datasets.
+
+        Parameters
+        ----------
+        datasets : List[Any]
+            List of datasets.
+        check_compatibility : bool, optional
+            Flag to check compatibility of datasets, by default True.
+        """
         self._check_compatibility = check_compatibility
         super().__init__(datasets)
 
-    def swap_with_parent(self, parent):
+    def swap_with_parent(self, parent: Any) -> Any:
+        """Swap datasets with the parent.
+
+        Parameters
+        ----------
+        parent : Any
+            Parent dataset.
+
+        Returns
+        -------
+        Any
+            New parent dataset with swapped datasets.
+        """
         new_parents = [parent.clone(ds) for ds in self.datasets]
         return self.clone(new_parents)
 
-    def clone(self, datasets):
+    def clone(self, datasets: List[Any]) -> "ZipBase":
+        """Clone the ZipBase with new datasets.
+
+        Parameters
+        ----------
+        datasets : List[Any]
+            List of new datasets.
+
+        Returns
+        -------
+        ZipBase
+            Cloned ZipBase instance.
+        """
         return self.__class__(datasets, check_compatibility=self._check_compatibility)
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the datasets.
+
+        Returns
+        -------
+        Node
+            Tree representation of the datasets.
+        """
         return Node(self, [d.tree() for d in self.datasets], check_compatibility=self._check_compatibility)
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Get the length of the smallest dataset.
+
+        Returns
+        -------
+        int
+            Length of the smallest dataset.
+        """
         return min(len(d) for d in self.datasets)
 
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> Tuple[Any, ...]:
+        """Get the item at the specified index from all datasets.
+
+        Parameters
+        ----------
+        n : FullIndex
+            Index to retrieve.
+
+        Returns
+        -------
+        Tuple[Any, ...]
+            Tuple of items from all datasets.
+        """
         return tuple(d[n] for d in self.datasets)
 
-    def check_same_resolution(self, d1, d2):
+    def check_same_resolution(self, d1: Dataset, d2: Dataset) -> None:
+        """Check if two datasets have the same resolution.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            First dataset.
+        d2 : Dataset
+            Second dataset.
+        """
         pass
 
-    def check_same_grid(self, d1, d2):
+    def check_same_grid(self, d1: Dataset, d2: Dataset) -> None:
+        """Check if two datasets have the same grid.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            First dataset.
+        d2 : Dataset
+            Second dataset.
+        """
         pass
 
-    def check_same_variables(self, d1, d2):
+    def check_same_variables(self, d1: Dataset, d2: Dataset) -> None:
+        """Check if two datasets have the same variables.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            First dataset.
+        d2 : Dataset
+            Second dataset.
+        """
         pass
 
     @cached_property
-    def missing(self):
-        result = set()
+    def missing(self) -> Set[int]:
+        """Get the set of missing indices from all datasets."""
+        result: Set[int] = set()
         for d in self.datasets:
             result = result | d.missing
         return result
 
     @property
-    def shape(self):
+    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):
+    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):
+    def latitudes(self) -> Tuple[Any, ...]:
+        """Get the latitudes of all datasets."""
         return tuple(d.latitudes for d in self.datasets)
 
     @property
-    def longitudes(self):
+    def longitudes(self) -> Tuple[Any, ...]:
+        """Get the longitudes of all datasets."""
         return tuple(d.longitudes for d in self.datasets)
 
     @property
-    def dtype(self):
+    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):
+    def grids(self) -> Tuple[Any, ...]:
+        """Get the grids of all datasets."""
         return tuple(d.grids for d in self.datasets)
 
     @property
-    def statistics(self):
+    def statistics(self) -> Tuple[Any, ...]:
+        """Get the statistics of all datasets."""
         return tuple(d.statistics for d in self.datasets)
 
     @property
-    def resolution(self):
+    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):
+    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)
 
-    def check_compatibility(self, d1, d2):
+    def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
+        """Check compatibility between two datasets.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            First dataset.
+        d2 : Dataset
+            Second dataset.
+        """
         if self._check_compatibility:
             super().check_compatibility(d1, d2)
 
 
 class Zip(ZipBase):
+    """Class for handling zipped datasets."""
+
     pass
 
 
 class XY(ZipBase):
+    """Class for handling XY datasets."""
+
     pass
 
 
-def xy_factory(args, kwargs):
+def xy_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> XY:
+    """Factory function to create an XY instance.
 
+    Parameters
+    ----------
+    args : Tuple[Any, ...]
+        Positional arguments.
+    kwargs : Dict[str, Any]
+        Keyword arguments.
+
+    Returns
+    -------
+    XY
+        An instance of XY.
+    """
     if "xy" in kwargs:
         xy = kwargs.pop("xy")
     else:
@@ -126,8 +256,21 @@ def xy_factory(args, kwargs):
     return XY(datasets, check_compatibility=check_compatibility)._subset(**kwargs)
 
 
-def zip_factory(args, kwargs):
+def zip_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Zip:
+    """Factory function to create a Zip instance.
+
+    Parameters
+    ----------
+    args : Tuple[Any, ...]
+        Positional arguments.
+    kwargs : Dict[str, Any]
+        Keyword arguments.
 
+    Returns
+    -------
+    Zip
+        An instance of Zip.
+    """
     zip = kwargs.pop("zip")
     assert len(args) == 0
     assert isinstance(zip, (list, tuple))

anemoi/datasets/dates/__init__.py

@@ -12,6 +12,13 @@ import datetime
 import warnings
 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
@@ -22,7 +29,17 @@ from anemoi.utils.hindcasts import HindcastDatesTimes
 from anemoi.utils.humanize import print_dates
 
 
-def extend(x):
+def extend(x: Union[str, List[Any], Tuple[Any, ...]]) -> Iterator[datetime.datetime]:
+    """Extend a date range or list of dates into individual datetime objects.
+
+    Args:
+        x (Union[str, List[Any], Tuple[Any, ...]]): A date range string or list/tuple of dates.
+
+    Returns
+    -------
+    Iterator[datetime.datetime]
+        An iterator of datetime objects.
+    """
 
     if isinstance(x, (list, tuple)):
         for y in x:
@@ -69,7 +86,14 @@ class DatesProvider:
     3
     """
 
-    def __init__(self, missing=None):
+    def __init__(self, missing: Optional[List[Union[str, datetime.datetime]]] = None) -> None:
+        """Initialize the DatesProvider with optional missing dates.
+
+        Parameters
+        ----------
+        missing : Optional[List[Union[str, datetime.datetime]]]
+            List of missing dates.
+        """
         if not missing:
             missing = []
         self.missing = list(extend(missing))
@@ -78,8 +102,17 @@ class DatesProvider:
             warnings.warn(f"Missing dates {len(diff)=} not in list.")
 
     @classmethod
-    def from_config(cls, **kwargs):
+    def from_config(cls, **kwargs: Any) -> "DatesProvider":
+        """Create a DatesProvider instance from configuration.
+
+        Args:
+            **kwargs (Any): Configuration parameters.
 
+        Returns
+        -------
+        DatesProvider
+            An instance of DatesProvider.
+        """
         if kwargs.pop("hindcasts", False):
             return HindcastsDates(**kwargs)
 
@@ -88,35 +121,113 @@ class DatesProvider:
 
         return StartEndDates(**kwargs)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[datetime.datetime]:
+        """Iterate over the dates.
+
+        Yields
+        ------
+        Iterator[datetime.datetime]
+            An iterator of datetime objects.
+        """
         yield from self.values
 
-    def __getitem__(self, i):
+    def __getitem__(self, i: int) -> datetime.datetime:
+        """Get a date by index.
+
+        Args:
+            i (int): Index of the date.
+
+        Returns
+        -------
+        datetime.datetime
+            The date at the specified index.
+        """
         return self.values[i]
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Get the number of dates.
+
+        Returns
+        -------
+        int
+            The number of dates.
+        """
         return len(self.values)
 
     @property
-    def summary(self):
+    def summary(self) -> str:
+        """Get a summary of the date range."""
         return f"📅 {self.values[0]} ... {self.values[-1]}"
 
 
 class ValuesDates(DatesProvider):
-    def __init__(self, values, **kwargs):
+    """Class for handling a list of date values.
+
+    Args:
+        values (List[Union[str, datetime.datetime]]): List of date values.
+        **kwargs (Any): Additional arguments.
+    """
+
+    def __init__(self, values: List[Union[str, datetime.datetime]], **kwargs: Any) -> None:
+        """Initialize ValuesDates with a list of values.
+
+        Args:
+            values (List[Union[str, datetime.datetime]]): List of date values.
+            **kwargs (Any): Additional arguments.
+        """
         self.values = sorted([as_datetime(_) for _ in values])
         super().__init__(**kwargs)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Get a string representation of the ValuesDates instance.
+
+        Returns
+        -------
+        str
+            String representation of the instance.
+        """
         return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
 
-    def as_dict(self):
+    def as_dict(self) -> Dict[str, Any]:
+        """Convert the ValuesDates instance to a dictionary.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Dictionary representation of the instance.
+        """
         return {"values": self.values[0]}
 
 
 class StartEndDates(DatesProvider):
-    def __init__(self, start, end, frequency=1, **kwargs):
+    """Class for generating dates between a start and end date with a specified frequency.
 
+    Args:
+        start (Union[str, datetime.datetime]): Start date.
+        end (Union[str, datetime.datetime]): End date.
+        frequency (Union[int, str]): Frequency of dates.
+        **kwargs (Any): Additional arguments.
+    """
+
+    def __repr__(self) -> str:
+        """Get a string representation of the StartEndDates instance."""
+        return f"{self.__class__.__name__}({self.start}..{self.end} every {self.frequency})"
+
+    def __init__(
+        self,
+        start: Union[str, datetime.datetime],
+        end: Union[str, datetime.datetime],
+        frequency: Union[int, str] = 1,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize StartEndDates with start, end, and frequency.
+
+        Args:
+            start (Union[str, datetime.datetime]): Start date.
+            end (Union[str, datetime.datetime]): End date.
+            frequency (Union[int, str]): Frequency of dates.
+            **kwargs (Any): Additional arguments.
+        """
         frequency = frequency_to_timedelta(frequency)
         assert isinstance(frequency, datetime.timedelta), frequency
 
@@ -148,7 +259,14 @@ class StartEndDates(DatesProvider):
 
         super().__init__(missing=missing)
 
-    def as_dict(self):
+    def as_dict(self) -> Dict[str, Any]:
+        """Convert the StartEndDates instance to a dictionary.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Dictionary representation of the instance.
+        """
         return {
             "start": self.start.isoformat(),
             "end": self.end.isoformat(),
@@ -157,8 +275,26 @@ class StartEndDates(DatesProvider):
 
 
 class Hindcast:
+    """Class representing a single hindcast date.
 
-    def __init__(self, date, refdate, hdate, step):
+    Args:
+        date (datetime.datetime): The date of the hindcast.
+        refdate (datetime.datetime): The reference date.
+        hdate (datetime.datetime): The hindcast date.
+        step (int): The step value.
+    """
+
+    def __init__(
+        self, date: datetime.datetime, refdate: datetime.datetime, hdate: datetime.datetime, step: int
+    ) -> None:
+        """Initialize a Hindcast instance.
+
+        Args:
+            date (datetime.datetime): The date of the hindcast.
+            refdate (datetime.datetime): The reference date.
+            hdate (datetime.datetime): The hindcast date.
+            step (int): The step value.
+        """
         self.date = date
         self.refdate = refdate
         self.hdate = hdate
@@ -166,8 +302,33 @@ class Hindcast:
 
 
 class HindcastsDates(DatesProvider):
-    def __init__(self, start, end, steps=[0], years=20, **kwargs):
+    """Class for generating hindcast dates over a range of years.
+
+    Args:
+        start (Union[str, List[str]]): Start date(s).
+        end (Union[str, List[str]]): End date(s).
+        steps (List[int]): List of step values.
+        years (int): Number of years.
+        **kwargs (Any): Additional arguments.
+    """
 
+    def __init__(
+        self,
+        start: Union[str, List[str]],
+        end: Union[str, List[str]],
+        steps: List[int] = [0],
+        years: int = 20,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize HindcastsDates with start, end, steps, and years.
+
+        Args:
+            start (Union[str, List[str]]): Start date(s).
+            end (Union[str, List[str]]): End date(s).
+            steps (List[int]): List of step values.
+            years (int): Number of years.
+            **kwargs (Any): Additional arguments.
+        """
         if not isinstance(start, list):
             start = [start]
             end = [end]
@@ -232,10 +393,24 @@ class HindcastsDates(DatesProvider):
 
         super().__init__(missing=missing)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Get a string representation of the HindcastsDates instance.
+
+        Returns
+        -------
+        str
+            String representation of the instance.
+        """
         return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
 
-    def as_dict(self):
+    def as_dict(self) -> Dict[str, Any]:
+        """Convert the HindcastsDates instance to a dictionary.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Dictionary representation of the instance.
+        """
         return {"hindcasts": self.hindcasts}