anemoi-utils 0.4.11__py3-none-any.whl → 0.4.13__py3-none-any.whl

anemoi/utils/__init__.py

@@ -7,6 +7,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+"""Anemoi Utils package."""
 
 try:
     # NOTE: the `_version.py` file must not be present in the git repository

anemoi/utils/__main__.py

@@ -7,6 +7,8 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+from typing import Any
+
 from anemoi.utils.cli import cli_main
 from anemoi.utils.cli import make_parser
 
@@ -15,11 +17,19 @@ from .commands import COMMANDS
 
 
 # For read-the-docs
-def create_parser():
+def create_parser() -> Any:
+    """Create the argument parser for the CLI.
+
+    Returns
+    -------
+    Any
+        The argument parser
+    """
     return make_parser(__doc__, COMMANDS)
 
 
-def main():
+def main() -> None:
+    """Main entry point for the CLI."""
     cli_main(__version__, __doc__, COMMANDS)
 
 

anemoi/utils/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.11'
-__version_tuple__ = version_tuple = (0, 4, 11)
+__version__ = version = '0.4.13'
+__version_tuple__ = version_tuple = (0, 4, 13)

anemoi/utils/caching.py

@@ -13,6 +13,9 @@ import json
 import os
 import time
 from threading import Lock
+from typing import Any
+from typing import Callable
+from typing import Optional
 
 import numpy as np
 
@@ -20,11 +23,30 @@ LOCK = Lock()
 CACHE = {}
 
 
-def _get_cache_path(collection):
+def _get_cache_path(collection: str) -> str:
+    """Get the cache path for a collection.
+
+    Parameters
+    ----------
+    collection : str
+        The name of the collection
+
+    Returns
+    -------
+    str
+        The cache path
+    """
     return os.path.join(os.path.expanduser("~"), ".cache", "anemoi", collection)
 
 
-def clean_cache(collection="default"):
+def clean_cache(collection: str = "default") -> None:
+    """Clean the cache for a collection.
+
+    Parameters
+    ----------
+    collection : str, optional
+        The name of the collection, by default "default"
+    """
     global CACHE
     CACHE = {}
     path = _get_cache_path(collection)
@@ -36,14 +58,35 @@ def clean_cache(collection="default"):
 
 class Cacher:
     """This class implements a simple caching mechanism.
-    Private class, do not use directly"""
+    Private class, do not use directly.
+    """
 
-    def __init__(self, collection, expires):
+    def __init__(self, collection: str, expires: Optional[int]):
+        """Initialize the Cacher.
+
+        Parameters
+        ----------
+        collection : str
+            The name of the collection
+        expires : int, optional
+            The expiration time in seconds, or None for no expiration
+        """
         self.collection = collection
         self.expires = expires
 
-    def __call__(self, func):
+    def __call__(self, func: Callable) -> Callable:
+        """Wrap a function with caching.
+
+        Parameters
+        ----------
+        func : Callable
+            The function to wrap
 
+        Returns
+        -------
+        Callable
+            The wrapped function
+        """
         full = f"{func.__module__}.{func.__name__}"
 
         def wrapped(*args, **kwargs):
@@ -55,8 +98,21 @@ class Cacher:
 
         return wrapped
 
-    def cache(self, key, proc):
-
+    def cache(self, key: tuple, proc: Callable) -> Any:
+        """Cache the result of a function.
+
+        Parameters
+        ----------
+        key : tuple
+            The cache key
+        proc : Callable
+            The function to call if the result is not cached
+
+        Returns
+        -------
+        Any
+            The cached result
+        """
         key = json.dumps(key, sort_keys=True)
         m = hashlib.md5()
         m.update(key.encode("utf-8"))
@@ -88,37 +144,106 @@ class Cacher:
 
 
 class JsonCacher(Cacher):
+    """Cacher that uses JSON files."""
+
     ext = ""
 
-    def save(self, path, data):
+    def save(self, path: str, data: dict) -> str:
+        """Save data to a JSON file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the JSON file
+        data : dict
+            The data to save
+
+        Returns
+        -------
+        str
+            The temporary file path
+        """
         temp_path = path + ".tmp"
         with open(temp_path, "w") as f:
             json.dump(data, f)
         return temp_path
 
-    def load(self, path):
+    def load(self, path: str) -> dict:
+        """Load data from a JSON file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the JSON file
+
+        Returns
+        -------
+        dict
+            The loaded data
+        """
         with open(path, "r") as f:
             return json.load(f)
 
 
 class NpzCacher(Cacher):
+    """Cacher that uses NPZ files."""
+
     ext = ".npz"
 
-    def save(self, path, data):
+    def save(self, path: str, data: dict) -> str:
+        """Save data to an NPZ file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the NPZ file
+        data : dict
+            The data to save
+
+        Returns
+        -------
+        str
+            The temporary file path
+        """
         temp_path = path + ".tmp.npz"
         np.savez(temp_path, **data)
         return temp_path
 
-    def load(self, path):
+    def load(self, path: str) -> dict:
+        """Load data from an NPZ file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the NPZ file
+
+        Returns
+        -------
+        dict
+            The loaded data
+        """
         return np.load(path, allow_pickle=True)
 
 
-# PUBLIC API
-def cached(collection="default", expires=None, encoding="json"):
+# This function is the main entry point for the caching mechanism for the other anemoi packages
+def cached(collection: str = "default", expires: Optional[int] = None, encoding: str = "json") -> Callable:
     """Decorator to cache the result of a function.
 
     Default is to use a json file to store the cache, but you can also use npz files
     to cache dict of numpy arrays.
 
+    Parameters
+    ----------
+    collection : str, optional
+        The name of the collection, by default "default"
+    expires : int, optional
+        The expiration time in seconds, or None for no expiration, by default None
+    encoding : str, optional
+        The encoding type, either "json" or "npz", by default "json"
+
+    Returns
+    -------
+    Callable
+        The decorated function
     """
     return dict(json=JsonCacher, npz=NpzCacher)[encoding](collection, expires)

anemoi/utils/checkpoints.py

@@ -18,6 +18,7 @@ import os
 import time
 import zipfile
 from tempfile import TemporaryDirectory
+from typing import Callable
 
 import tqdm
 
@@ -28,7 +29,7 @@ DEFAULT_FOLDER = "anemoi-metadata"
 
 
 def has_metadata(path: str, *, name: str = DEFAULT_NAME) -> bool:
-    """Check if a checkpoint file has a metadata file
+    """Check if a checkpoint file has a metadata file.
 
     Parameters
     ----------
@@ -49,8 +50,26 @@ def has_metadata(path: str, *, name: str = DEFAULT_NAME) -> bool:
     return False
 
 
-def metadata_root(path: str, *, name: str = DEFAULT_NAME) -> bool:
+def metadata_root(path: str, *, name: str = DEFAULT_NAME) -> str:
+    """Get the root directory of the metadata file.
 
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    name : str, optional
+        The name of the metadata file in the zip archive
+
+    Returns
+    -------
+    str
+        The root directory of the metadata file
+
+    Raises
+    ------
+    ValueError
+        If the metadata file is not found
+    """
     with zipfile.ZipFile(path, "r") as f:
         for b in f.namelist():
             if os.path.basename(b) == name:
@@ -58,15 +77,15 @@ def metadata_root(path: str, *, name: str = DEFAULT_NAME) -> bool:
     raise ValueError(f"Could not find '{name}' in {path}.")
 
 
-def load_metadata(path: str, *, supporting_arrays=False, name: str = DEFAULT_NAME) -> dict:
-    """Load metadata from a checkpoint file
+def load_metadata(path: str, *, supporting_arrays: bool = False, name: str = DEFAULT_NAME) -> dict:
+    """Load metadata from a checkpoint file.
 
     Parameters
     ----------
     path : str
         The path to the checkpoint file
 
-    supporting_arrays: bool, optional
+    supporting_arrays : bool, optional
         If True, the function will return a dictionary with the supporting arrays
 
     name : str, optional
@@ -102,7 +121,21 @@ def load_metadata(path: str, *, supporting_arrays=False, name: str = DEFAULT_NAM
         raise ValueError(f"Could not find '{name}' in {path}.")
 
 
-def load_supporting_arrays(zipf, entries) -> dict:
+def load_supporting_arrays(zipf: zipfile.ZipFile, entries: dict) -> dict:
+    """Load supporting arrays from a zip file.
+
+    Parameters
+    ----------
+    zipf : zipfile.ZipFile
+        The zip file
+    entries : dict
+        A dictionary of entries with paths, shapes, and dtypes
+
+    Returns
+    -------
+    dict
+        A dictionary of supporting arrays
+    """
     import numpy as np
 
     supporting_arrays = {}
@@ -114,16 +147,18 @@ def load_supporting_arrays(zipf, entries) -> dict:
     return supporting_arrays
 
 
-def save_metadata(path, metadata, *, supporting_arrays=None, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> None:
-    """Save metadata to a checkpoint file
+def save_metadata(
+    path: str, metadata: dict, *, supporting_arrays: dict = None, name: str = DEFAULT_NAME, folder: str = DEFAULT_FOLDER
+) -> None:
+    """Save metadata to a checkpoint file.
 
     Parameters
     ----------
     path : str
         The path to the checkpoint file
-    metadata : JSON
+    metadata : dict
         A JSON serializable object
-    supporting_arrays: dict, optional
+    supporting_arrays : dict, optional
         A dictionary of supporting NumPy arrays
     name : str, optional
         The name of the metadata file in the zip archive
@@ -179,7 +214,20 @@ def save_metadata(path, metadata, *, supporting_arrays=None, name=DEFAULT_NAME,
             zipf.writestr(entry["path"], value.tobytes())
 
 
-def _edit_metadata(path, name, callback, supporting_arrays=None):
+def _edit_metadata(path: str, name: str, callback: Callable, supporting_arrays: dict = None) -> None:
+    """Edit metadata in a checkpoint file.
+
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    name : str
+        The name of the metadata file in the zip archive
+    callback : Callable
+        A callback function to edit the metadata
+    supporting_arrays : dict, optional
+        A dictionary of supporting NumPy arrays
+    """
     new_path = f"{path}.anemoi-edit-{time.time()}-{os.getpid()}.tmp"
 
     found = False
@@ -223,8 +271,20 @@ def _edit_metadata(path, name, callback, supporting_arrays=None):
     LOG.info("Updated metadata in %s", path)
 
 
-def replace_metadata(path, metadata, supporting_arrays=None, *, name=DEFAULT_NAME):
+def replace_metadata(path: str, metadata: dict, supporting_arrays: dict = None, *, name: str = DEFAULT_NAME) -> None:
+    """Replace metadata in a checkpoint file.
 
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    metadata : dict
+        A JSON serializable object
+    supporting_arrays : dict, optional
+        A dictionary of supporting NumPy arrays
+    name : str, optional
+        The name of the metadata file in the zip archive
+    """
     if not isinstance(metadata, dict):
         raise ValueError(f"metadata must be a dict, got {type(metadata)}")
 
@@ -238,8 +298,16 @@ def replace_metadata(path, metadata, supporting_arrays=None, *, name=DEFAULT_NAM
     return _edit_metadata(path, name, callback, supporting_arrays)
 
 
-def remove_metadata(path, *, name=DEFAULT_NAME):
+def remove_metadata(path: str, *, name: str = DEFAULT_NAME) -> None:
+    """Remove metadata from a checkpoint file.
 
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    name : str, optional
+        The name of the metadata file in the zip archive
+    """
     LOG.info("Removing metadata '%s' from %s", name, path)
 
     def callback(full):

anemoi/utils/cli.py

@@ -14,6 +14,7 @@ import logging
 import os
 import sys
 import traceback
+from typing import Callable
 
 try:
     import argcomplete
@@ -24,13 +25,36 @@ LOG = logging.getLogger(__name__)
 
 
 class Command:
+    """Base class for commands."""
+
     accept_unknown_args = False
 
-    def run(self, args):
+    def run(self, args: argparse.Namespace) -> None:
+        """Run the command.
+
+        Parameters
+        ----------
+        args : argparse.Namespace
+            The arguments for the command
+        """
         raise NotImplementedError(f"Command not implemented: {args.command}")
 
 
-def make_parser(description, commands):
+def make_parser(description: str, commands: dict[str, Command]) -> argparse.ArgumentParser:
+    """Create an argument parser for the CLI.
+
+    Parameters
+    ----------
+    description : str
+        The description of the CLI
+    commands : dict[str, Command]
+        A dictionary of command names to Command instances
+
+    Returns
+    -------
+    argparse.ArgumentParser
+        The argument parser
+    """
     parser = argparse.ArgumentParser(
         description=description,
         formatter_class=argparse.RawDescriptionHelpFormatter,
@@ -60,20 +84,61 @@ def make_parser(description, commands):
 class Failed(Command):
     """Command not available."""
 
-    def __init__(self, name, error):
+    def __init__(self, name: str, error: ImportError):
+        """Initialize the Failed command.
+
+        Parameters
+        ----------
+        name : str
+            The name of the command
+        error : ImportError
+            The error that occurred
+        """
         self.name = name
         self.error = error
         traceback.print_tb(error.__traceback__)
 
-    def add_arguments(self, command_parser):
+    def add_arguments(self, command_parser: argparse.ArgumentParser) -> None:
+        """Add arguments to the command parser.
+
+        Parameters
+        ----------
+        command_parser : argparse.ArgumentParser
+            The command parser
+        """
         command_parser.add_argument("x", nargs=argparse.REMAINDER)
 
-    def run(self, args):
+    def run(self, args: argparse.Namespace) -> None:
+        """Run the command.
+
+        Parameters
+        ----------
+        args : argparse.Namespace
+            The arguments for the command
+        """
         print(f"Command '{self.name}' not available: {self.error}")
         sys.exit(1)
 
 
-def register_commands(here, package, select, fail=None):
+def register_commands(here: str, package: str, select: Callable, fail: Callable = None) -> dict[str, Command]:
+    """Register commands from a package.
+
+    Parameters
+    ----------
+    here : str
+        The directory containing the commands
+    package : str
+        The package name
+    select : Callable
+        A function to select the command object from the module
+    fail : Callable, optional
+        A function to create a Failed command if a command cannot be imported
+
+    Returns
+    -------
+    dict[str, Command]
+        A dictionary of command names to Command instances
+    """
     result = {}
     not_available = {}
 
@@ -120,7 +185,18 @@ def register_commands(here, package, select, fail=None):
     return result
 
 
-def cli_main(version, description, commands):
+def cli_main(version: str, description: str, commands: dict[str, Command]) -> None:
+    """Main entry point for the CLI.
+
+    Parameters
+    ----------
+    version : str
+        The version of the CLI
+    description : str
+        The description of the CLI
+    commands : dict[str, Command]
+        A dictionary of command names to Command instances
+    """
     parser = make_parser(description, commands)
     args, unknown = parser.parse_known_args()
     if argcomplete:

anemoi/utils/commands/__init__.py

@@ -7,6 +7,10 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+"""This module initializes and registers command-line interface (CLI) commands
+for the Anemoi utilities.
+"""
+
 import os
 
 from anemoi.utils.cli import Command

anemoi/utils/commands/config.py

@@ -9,6 +9,8 @@
 
 
 import json
+from argparse import ArgumentParser
+from argparse import Namespace
 
 from ..config import config_path
 from ..config import load_config
@@ -16,11 +18,26 @@ from . import Command
 
 
 class Config(Command):
+    """Handle configuration related commands."""
 
-    def add_arguments(self, command_parser):
+    def add_arguments(self, command_parser: ArgumentParser) -> None:
+        """Add arguments to the command parser.
+
+        Parameters
+        ----------
+        command_parser : ArgumentParser
+            The argument parser to which the arguments will be added.
+        """
         command_parser.add_argument("--path", help="Print path to config file")
 
-    def run(self, args):
+    def run(self, args: Namespace) -> None:
+        """Execute the command with the provided arguments.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
         if args.path:
             print(config_path())
         else:

anemoi/utils/commands/requests.py

@@ -6,6 +6,9 @@
 # nor does it submit to any jurisdiction.
 
 import json
+import sys
+from argparse import ArgumentParser
+from argparse import Namespace
 
 from anemoi.utils.mars.requests import print_request
 
@@ -15,15 +18,32 @@ from . import Command
 class Requests(Command):
     """Convert a JSON requests file to MARS format."""
 
-    def add_arguments(self, command_parser):
+    def add_arguments(self, command_parser: ArgumentParser) -> None:
+        """Add arguments to the command parser.
+
+        Parameters
+        ----------
+        command_parser : ArgumentParser
+            The argument parser to which the arguments will be added.
+        """
         command_parser.add_argument("input")
         command_parser.add_argument("output")
         command_parser.add_argument("--verb", default="retrieve")
         command_parser.add_argument("--only-one-field", action="store_true")
 
-    def run(self, args):
-        with open(args.input) as f:
-            requests = json.load(f)
+    def run(self, args: Namespace) -> None:
+        """Execute the command with the provided arguments.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        if args.input == "-":
+            requests = json.load(sys.stdin)
+        else:
+            with open(args.input) as f:
+                requests = json.load(f)
 
         if args.only_one_field:
             for r in requests: