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

anemoi/datasets/create/persistent.py

@@ -16,6 +16,9 @@ import os
 import pickle
 import shutil
 import socket
+from typing import Any
+from typing import Iterator
+from typing import Tuple
 
 import numpy as np
 from anemoi.utils.provenance import gather_provenance_info
@@ -24,30 +27,62 @@ LOG = logging.getLogger(__name__)
 
 
 class PersistentDict:
+    """A dictionary-like object that persists its contents to disk using pickle files.
+
+    Attributes
+    ----------
+    version : int
+        The version of the PersistentDict.
+    dirname : str
+        The directory where the data is stored.
+    name : str
+        The name of the directory.
+    ext : str
+        The extension of the directory.
+    """
+
     version = 3
 
     # Used in parrallel, during data loading,
     # to write data in pickle files.
-    def __init__(self, directory, create=True):
-        """dirname: str The directory where the data will be stored."""
+    def __init__(self, directory: str, create: bool = True):
+        """Initialize the PersistentDict.
+
+        Parameters
+        ----------
+        directory : str
+            The directory where the data will be stored.
+        create : bool, optional
+            Whether to create the directory if it doesn't exist.
+        """
         self.dirname = directory
         self.name, self.ext = os.path.splitext(os.path.basename(self.dirname))
         if create:
             self.create()
 
-    def create(self):
+    def create(self) -> None:
+        """Create the directory if it doesn't exist."""
         os.makedirs(self.dirname, exist_ok=True)
 
-    def delete(self):
+    def delete(self) -> None:
+        """Delete the directory and its contents."""
         try:
             shutil.rmtree(self.dirname)
         except FileNotFoundError:
             pass
 
-    def __str__(self):
+    def __str__(self) -> str:
+        """Return a string representation of the PersistentDict."""
         return f"{self.__class__.__name__}({self.dirname})"
 
-    def items(self):
+    def items(self) -> Iterator[Any]:
+        """Yield items stored in the directory.
+
+        Yields
+        ------
+        Iterator[Any]
+            An iterator over the items.
+        """
         # use glob to read all pickles
         files = glob.glob(self.dirname + "/*.pickle")
         LOG.debug(f"Reading {self.name} data, found {len(files)} files in {self.dirname}")
@@ -56,7 +91,14 @@ class PersistentDict:
             with open(f, "rb") as f:
                 yield pickle.load(f)
 
-    def add_provenance(self, **kwargs):
+    def add_provenance(self, **kwargs: Any) -> None:
+        """Add provenance information to the directory.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional provenance information.
+        """
         path = os.path.join(self.dirname, "provenance.json")
         if os.path.exists(path):
             return
@@ -64,10 +106,28 @@ class PersistentDict:
         with open(path, "w") as f:
             json.dump(out, f)
 
-    def add(self, elt, *, key):
+    def add(self, elt: Any, *, key: Any) -> None:
+        """Add an element to the PersistentDict.
+
+        Parameters
+        ----------
+        elt : Any
+            The element to add.
+        key : Any
+            The key associated with the element.
+        """
         self[key] = elt
 
-    def __setitem__(self, key, elt):
+    def __setitem__(self, key: Any, elt: Any) -> None:
+        """Set an item in the PersistentDict.
+
+        Parameters
+        ----------
+        key : Any
+            The key associated with the element.
+        elt : Any
+            The element to set.
+        """
         h = hashlib.sha256(str(key).encode("utf-8")).hexdigest()
         path = os.path.join(self.dirname, f"{h}.pickle")
 
@@ -81,42 +141,99 @@ class PersistentDict:
 
         LOG.debug(f"Written {self.name} data for len {key} in {path}")
 
-    def flush(self):
+    def flush(self) -> None:
+        """Flush the PersistentDict (no-op)."""
         pass
 
 
 class BufferedPersistentDict(PersistentDict):
-    def __init__(self, buffer_size=1000, **kwargs):
+    """A buffered version of PersistentDict that stores elements in memory before persisting them to disk.
+
+    Attributes
+    ----------
+    buffer_size : int
+        The size of the buffer.
+    elements : list
+        The list of elements in the buffer.
+    keys : list
+        The list of keys in the buffer.
+    storage : PersistentDict
+        The underlying PersistentDict used for storage.
+    """
+
+    def __init__(self, buffer_size: int = 1000, **kwargs: Any):
+        """Initialize the BufferedPersistentDict.
+
+        Parameters
+        ----------
+        buffer_size : int, optional
+            The size of the buffer.
+        **kwargs : Any
+            Additional arguments for PersistentDict.
+        """
         self.buffer_size = buffer_size
         self.elements = []
         self.keys = []
         self.storage = PersistentDict(**kwargs)
 
-    def add(self, elt, *, key):
+    def add(self, elt: Any, *, key: Any) -> None:
+        """Add an element to the BufferedPersistentDict.
+
+        Parameters
+        ----------
+        elt : Any
+            The element to add.
+        key : Any
+            The key associated with the element.
+        """
         self.elements.append(elt)
         self.keys.append(key)
         if len(self.keys) > self.buffer_size:
             self.flush()
 
-    def flush(self):
+    def flush(self) -> None:
+        """Flush the buffer and store the elements in PersistentDict."""
         k = sorted(self.keys)
         self.storage.add(self.elements, key=k)
         self.elements = []
         self.keys = []
 
-    def items(self):
+    def items(self) -> Iterator[Tuple[Any, Any]]:
+        """Yield items stored in the BufferedPersistentDict.
+
+        Yields
+        ------
+        Iterator[Tuple[Any, Any]]
+            An iterator over the items.
+        """
         for keys, elements in self.storage.items():
             for key, elt in zip(keys, elements):
                 yield key, elt
 
-    def delete(self):
+    def delete(self) -> None:
+        """Delete the storage directory and its contents."""
         self.storage.delete()
 
-    def create(self):
+    def create(self) -> None:
+        """Create the storage directory if it doesn't exist."""
         self.storage.create()
 
 
-def build_storage(directory, create=True):
+def build_storage(directory: str, create: bool = True) -> BufferedPersistentDict:
+    """Build a BufferedPersistentDict storage.
+
+    Parameters
+    ----------
+    directory : str
+        The directory where the data will be stored.
+    create : bool, optional
+        Whether to create the directory if it doesn't exist.
+
+    Returns
+    -------
+    BufferedPersistentDict
+        The created BufferedPersistentDict.
+    """
     return BufferedPersistentDict(directory=directory, create=create)
 
 

anemoi/datasets/create/size.py

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

anemoi/datasets/create/source.py

@@ -0,0 +1,51 @@
+# (C) Copyright 2025- Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from abc import ABC
+from abc import abstractmethod
+
+import earthkit.data as ekd
+
+from anemoi.datasets.create.typing import DateList
+
+
+class Source(ABC):
+    """Represents a data source with a given context."""
+
+    emoji = "📦"  # For tracing
+
+    def __init__(self, context: any, *args: tuple, **kwargs: dict):
+        """Initialise the source.
+        Parameters
+        ----------
+        context : Any
+            The context for the data source.
+        *args : tuple
+            Additional positional arguments.
+        **kwargs : dict
+            Additional keyword arguments.
+        """
+        self.context = context
+
+    @abstractmethod
+    def execute(self, dates: DateList) -> ekd.FieldList:
+        """Execute the filter.
+
+        Parameters
+        ----------
+        dates : DateList
+            The input dates.
+
+        Returns
+        -------
+        ekd.FieldList
+            The output data.
+        """
+
+        pass

anemoi/datasets/create/sources/__init__.py

@@ -0,0 +1,36 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import logging
+from typing import Any
+
+from anemoi.utils.registry import Registry
+
+LOG = logging.getLogger(__name__)
+
+
+source_registry = Registry(__name__)
+
+
+def create_source(context: Any, config: Any) -> Any:
+    """Create a source based on the provided configuration.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the source is created.
+    config : Any
+        The configuration for the source.
+
+    Returns
+    -------
+    Any
+        The created source.
+    """
+    return source_registry.from_config(config, context)