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

anemoi/datasets/create/utils.py

@@ -12,12 +12,28 @@ import datetime
 import os
 import warnings
 from contextlib import contextmanager
+from typing import Any
+from typing import Union
 
 import numpy as np
 from earthkit.data import settings
+from numpy.typing import NDArray
 
 
-def cache_context(dirname):
+def cache_context(dirname: str) -> contextmanager:
+    """Context manager for setting a temporary cache directory.
+
+    Parameters
+    ----------
+    dirname : str
+        The directory name for the cache.
+
+    Returns
+    -------
+    contextmanager
+        A context manager that sets the cache directory.
+    """
+
     @contextmanager
     def no_cache_context():
         yield
@@ -30,7 +46,21 @@ def cache_context(dirname):
     return settings.temporary({"cache-policy": "user", "user-cache-directory": dirname})
 
 
-def to_datetime_list(*args, **kwargs):
+def to_datetime_list(*args: Any, **kwargs: Any) -> list[datetime.datetime]:
+    """Convert various date formats to a list of datetime objects.
+
+    Parameters
+    ----------
+    *args : Any
+        Positional arguments for date conversion.
+    **kwargs : Any
+        Keyword arguments for date conversion.
+
+    Returns
+    -------
+    list[datetime.datetime]
+        A list of datetime objects.
+    """
     from earthkit.data.utils.dates import to_datetime_list as to_datetime_list_
 
     warnings.warn(
@@ -41,7 +71,21 @@ def to_datetime_list(*args, **kwargs):
     return to_datetime_list_(*args, **kwargs)
 
 
-def to_datetime(*args, **kwargs):
+def to_datetime(*args: Any, **kwargs: Any) -> datetime.datetime:
+    """Convert various date formats to a single datetime object.
+
+    Parameters
+    ----------
+    *args : Any
+        Positional arguments for date conversion.
+    **kwargs : Any
+        Keyword arguments for date conversion.
+
+    Returns
+    -------
+    datetime.datetime
+        A datetime object.
+    """
     from earthkit.data.utils.dates import to_datetime as to_datetime_
 
     warnings.warn(
@@ -53,7 +97,24 @@ def to_datetime(*args, **kwargs):
     return to_datetime_(*args, **kwargs)
 
 
-def make_list_int(value):
+def make_list_int(value: Union[str, list, tuple, int]) -> list[int]:
+    """Convert a string, list, tuple, or integer to a list of integers.
+
+    Parameters
+    ----------
+    value : str or list or tuple or int
+        The value to convert.
+
+    Returns
+    -------
+    list[int]
+        A list of integers.
+
+    Raises
+    ------
+    ValueError
+        If the value cannot be converted to a list of integers.
+    """
     # Convert a string like "1/2/3" or "1/to/3" or "1/to/10/by/2" to a list of integers.
     # Moved to anemoi.utils.humanize
     # replace with from anemoi.utils.humanize import make_list_int
@@ -78,8 +139,38 @@ def make_list_int(value):
     raise ValueError(f"Cannot make list from {value}")
 
 
-def normalize_and_check_dates(dates, start, end, frequency, dtype="datetime64[s]"):
-
+def normalize_and_check_dates(
+    dates: list[datetime.datetime],
+    start: datetime.datetime,
+    end: datetime.datetime,
+    frequency: datetime.timedelta,
+    dtype: str = "datetime64[s]",
+) -> NDArray[Any]:
+    """Normalize and check a list of dates against a specified frequency.
+
+    Parameters
+    ----------
+    dates : list[datetime.datetime]
+        The list of dates to check.
+    start : datetime.datetime
+        The start date.
+    end : datetime.datetime
+        The end date.
+    frequency : datetime.timedelta
+        The frequency of the dates.
+    dtype : str, optional
+        The data type of the dates, by default "datetime64[s]".
+
+    Returns
+    -------
+    NDArray[Any]
+        An array of normalized dates.
+
+    Raises
+    ------
+    ValueError
+        If the final date size does not match the data shape.
+    """
     dates = [d.hdate if hasattr(d, "hdate") else d for d in dates]
 
     assert isinstance(frequency, datetime.timedelta), frequency

anemoi/datasets/create/writer.py

@@ -9,8 +9,10 @@
 
 
 import logging
+from typing import Any
 
 import numpy as np
+from numpy.typing import NDArray
 
 LOG = logging.getLogger(__name__)
 
@@ -23,20 +25,40 @@ class ViewCacheArray:
     temporarily store the data before flushing it to the array.
 
     The `flush` method copies the contents of the cache to the final array.
-
     """
 
-    def __init__(self, array, *, shape, indexes):
+    def __init__(self, array: NDArray[Any], *, shape: tuple[int, ...], indexes: list[int]):
+        """Initialize the ViewCacheArray.
+
+        Parameters
+        ----------
+        array : NDArray[Any]
+            The NumPy-like array to store the final data.
+        shape : tuple[int, ...]
+            The shape of the cache array.
+        indexes : list[int]
+            List to reindex the first dimension.
+        """
         assert len(indexes) == shape[0], (len(indexes), shape[0])
         self.array = array
         self.dtype = array.dtype
         self.cache = np.full(shape, np.nan, dtype=self.dtype)
         self.indexes = indexes
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: tuple[int, ...], value: NDArray[Any]) -> None:
+        """Set the value in the cache array at the specified key.
+
+        Parameters
+        ----------
+        key : tuple[int, ...]
+            The index key to set the value.
+        value : NDArray[Any]
+            The value to set in the cache array.
+        """
         self.cache[key] = value
 
-    def flush(self):
+    def flush(self) -> None:
+        """Copy the contents of the cache to the final array."""
         for i in range(self.cache.shape[0]):
             global_i = self.indexes[i]
             self.array[global_i] = self.cache[i]

anemoi/datasets/create/zarr.py

@@ -10,24 +10,56 @@
 import datetime
 import logging
 import shutil
+from typing import Any
+from typing import Optional
 
 import numpy as np
+import zarr
+from numpy.typing import NDArray
 
 LOG = logging.getLogger(__name__)
 
 
 def add_zarr_dataset(
     *,
-    name,
-    dtype=None,
-    fill_value=None,
-    zarr_root,
-    shape=None,
-    array=None,
-    overwrite=True,
-    dimensions=None,
+    name: str,
+    dtype: np.dtype = None,
+    fill_value: np.generic = None,
+    zarr_root: zarr.Group,
+    shape: tuple[int, ...] = None,
+    array: NDArray[Any] = None,
+    overwrite: bool = True,
+    dimensions: tuple[str, ...] = None,
     **kwargs,
-):
+) -> zarr.Array:
+    """Add a dataset to a Zarr group.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset.
+    dtype : np.dtype, optional
+        Data type of the dataset.
+    fill_value : np.generic, optional
+        Fill value for the dataset.
+    zarr_root : zarr.Group
+        Root Zarr group.
+    shape : tuple[int, ...], optional
+        Shape of the dataset.
+    array : NDArray[Any], optional
+        Array to initialize the dataset with.
+    overwrite : bool
+        Whether to overwrite existing dataset.
+    dimensions : tuple[str, ...]
+        Dimensions of the dataset.
+    **kwargs
+        Additional arguments for Zarr dataset creation.
+
+    Returns
+    -------
+    zarr.Array
+        The created Zarr array.
+    """
     assert dimensions is not None, "Please pass dimensions to add_zarr_dataset."
     assert isinstance(dimensions, (tuple, list))
 
@@ -80,13 +112,26 @@ def add_zarr_dataset(
 
 
 class ZarrBuiltRegistry:
+    """A class to manage the creation and access of Zarr datasets."""
+
     name_lengths = "lengths"
     name_flags = "flags"
     lengths = None
     flags = None
     z = None
 
-    def __init__(self, path, synchronizer_path=None, use_threads=False):
+    def __init__(self, path: str, synchronizer_path: Optional[str] = None, use_threads: bool = False):
+        """Initialize the ZarrBuiltRegistry.
+
+        Parameters
+        ----------
+        path : str
+            Path to the Zarr store.
+        synchronizer_path : Optional[str], optional
+            Path to the synchronizer.
+        use_threads : bool
+            Whether to use thread-based synchronization.
+        """
         import zarr
 
         assert isinstance(path, str), path
@@ -101,19 +146,33 @@ class ZarrBuiltRegistry:
             self.synchronizer_path = synchronizer_path
             self.synchronizer = zarr.ProcessSynchronizer(self.synchronizer_path)
 
-    def clean(self):
+    def clean(self) -> None:
+        """Clean up the synchronizer path."""
         if self.synchronizer_path is not None:
             try:
                 shutil.rmtree(self.synchronizer_path)
             except FileNotFoundError:
                 pass
 
-    def _open_write(self):
+    def _open_write(self) -> zarr.Group:
+        """Open the Zarr store in write mode."""
         import zarr
 
         return zarr.open(self.zarr_path, mode="r+", synchronizer=self.synchronizer)
 
-    def _open_read(self, sync=True):
+    def _open_read(self, sync: bool = True) -> zarr.Group:
+        """Open the Zarr store in read mode.
+
+        Parameters
+        ----------
+        sync : bool
+            Whether to use synchronization.
+
+        Returns
+        -------
+        zarr.Group
+            The opened Zarr group.
+        """
         import zarr
 
         if sync:
@@ -121,12 +180,30 @@ class ZarrBuiltRegistry:
         else:
             return zarr.open(self.zarr_path, mode="r")
 
-    def new_dataset(self, *args, **kwargs):
+    def new_dataset(self, *args, **kwargs) -> None:
+        """Create a new dataset in the Zarr store.
+
+        Parameters
+        ----------
+        *args
+            Positional arguments for dataset creation.
+        **kwargs
+            Keyword arguments for dataset creation.
+        """
         z = self._open_write()
         zarr_root = z["_build"]
         add_zarr_dataset(*args, zarr_root=zarr_root, overwrite=True, dimensions=("tmp",), **kwargs)
 
-    def add_to_history(self, action, **kwargs):
+    def add_to_history(self, action: str, **kwargs) -> None:
+        """Add an action to the history attribute of the Zarr store.
+
+        Parameters
+        ----------
+        action : str
+            The action to record.
+        **kwargs
+            Additional information about the action.
+        """
         new = dict(
             action=action,
             timestamp=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None).isoformat(),
@@ -138,37 +215,107 @@ class ZarrBuiltRegistry:
         history.append(new)
         z.attrs["history"] = history
 
-    def get_lengths(self):
+    def get_lengths(self) -> list[int]:
+        """Get the lengths dataset.
+
+        Returns
+        -------
+        list[int]
+            The lengths dataset.
+        """
         z = self._open_read()
         return list(z["_build"][self.name_lengths][:])
 
-    def get_flags(self, **kwargs):
+    def get_flags(self, **kwargs) -> list[bool]:
+        """Get the flags dataset.
+
+        Parameters
+        ----------
+        **kwargs
+            Additional arguments for reading the dataset.
+
+        Returns
+        -------
+        list[bool]
+            The flags dataset.
+        """
         z = self._open_read(**kwargs)
         return list(z["_build"][self.name_flags][:])
 
-    def get_flag(self, i):
+    def get_flag(self, i: int) -> bool:
+        """Get a specific flag.
+
+        Parameters
+        ----------
+        i : int
+            Index of the flag.
+
+        Returns
+        -------
+        bool
+            The flag value.
+        """
         z = self._open_read()
         return z["_build"][self.name_flags][i]
 
-    def set_flag(self, i, value=True):
+    def set_flag(self, i: int, value: bool = True) -> None:
+        """Set a specific flag.
+
+        Parameters
+        ----------
+        i : int
+            Index of the flag.
+        value : bool
+            Value to set the flag to.
+        """
         z = self._open_write()
         z.attrs["latest_write_timestamp"] = (
             datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None).isoformat()
         )
         z["_build"][self.name_flags][i] = value
 
-    def ready(self):
+    def ready(self) -> bool:
+        """Check if all flags are set.
+
+        Returns
+        -------
+        bool
+            True if all flags are set, False otherwise.
+        """
         return all(self.get_flags())
 
-    def create(self, lengths, overwrite=False):
+    def create(self, lengths: list[int], overwrite: bool = False) -> None:
+        """Create the lengths and flags datasets.
+
+        Parameters
+        ----------
+        lengths : list[int]
+            Lengths to initialize the dataset with.
+        overwrite : bool
+            Whether to overwrite existing datasets.
+        """
         self.new_dataset(name=self.name_lengths, array=np.array(lengths, dtype="i4"))
         self.new_dataset(name=self.name_flags, array=np.array([False] * len(lengths), dtype=bool))
         self.add_to_history("initialised")
 
-    def reset(self, lengths):
+    def reset(self, lengths: list[int]) -> None:
+        """Reset the lengths and flags datasets.
+
+        Parameters
+        ----------
+        lengths : list[int]
+            Lengths to initialize the dataset with.
+        """
         return self.create(lengths, overwrite=True)
 
-    def add_provenance(self, name):
+    def add_provenance(self, name: str) -> None:
+        """Add provenance information to the Zarr store.
+
+        Parameters
+        ----------
+        name : str
+            Name of the provenance attribute.
+        """
         z = self._open_write()
 
         if name in z.attrs:

anemoi/datasets/data/__init__.py

@@ -8,11 +8,20 @@
 # nor does it submit to any jurisdiction.
 
 import logging
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Set
 
+# from .dataset import FullIndex
+# from .dataset import Shape
+# from .dataset import TupleIndex
 from .misc import _open_dataset
 from .misc import add_dataset_path
 from .misc import add_named_dataset
 
+if TYPE_CHECKING:
+    from .dataset import Dataset
+
 LOG = logging.getLogger(__name__)
 
 __all__ = [
@@ -27,8 +36,19 @@ class MissingDateError(Exception):
     pass
 
 
-def _convert(x):
+def _convert(x: Any) -> Any:
+    """Convert OmegaConf objects to standard Python containers.
+
+    Parameters
+    ----------
+    x : Any
+        The object to convert.
 
+    Returns
+    -------
+    Any
+        The converted object.
+    """
     if isinstance(x, list):
         return [_convert(a) for a in x]
 
@@ -46,8 +66,21 @@ def _convert(x):
     return x
 
 
-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.
+    """
     # That will get rid of OmegaConf objects
 
     args, kwargs = _convert(args), _convert(kwargs)
@@ -59,8 +92,22 @@ def open_dataset(*args, **kwargs):
     return ds
 
 
-def list_dataset_names(*args, **kwargs):
+def list_dataset_names(*args: Any, **kwargs: Any) -> list[str]:
+    """List the names of datasets.
+
+    Parameters
+    ----------
+    *args : Any
+        Positional arguments.
+    **kwargs : Any
+        Keyword arguments.
+
+    Returns
+    -------
+    list of str
+        The list of dataset names.
+    """
     ds = _open_dataset(*args, **kwargs)
-    names = set()
+    names: Set[str] = set()
     ds.get_dataset_names(names)
     return sorted(names)