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

anemoi/utils/logs.py

@@ -19,17 +19,45 @@ thread_local = threading.local()
 LOGGER = logging.getLogger(__name__)
 
 
-def set_logging_name(name):
+def set_logging_name(name: str) -> None:
+    """Set the logging name for the current thread.
+
+    Parameters
+    ----------
+    name : str
+        The name to set for logging.
+    """
     thread_local.logging_name = name
 
 
 class ThreadCustomFormatter(logging.Formatter):
-    def format(self, record):
+    """Custom logging formatter that includes thread-specific logging names."""
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Format the log record to include the thread-specific logging name.
+
+        Parameters
+        ----------
+        record : logging.LogRecord
+            The log record to format.
+
+        Returns
+        -------
+        str
+            The formatted log record.
+        """
         record.logging_name = thread_local.logging_name
         return super().format(record)
 
 
-def enable_logging_name(name="main"):
+def enable_logging_name(name: str = "main") -> None:
+    """Enable logging with a thread-specific logging name.
+
+    Parameters
+    ----------
+    name : str, optional
+        The default logging name to set, by default "main".
+    """
     thread_local.logging_name = name
 
     formatter = ThreadCustomFormatter("%(asctime)s - %(logging_name)s - %(levelname)s - %(message)s")

anemoi/utils/mars/__init__.py

@@ -10,13 +10,16 @@
 
 """Utilities for working with Mars requests.
 
-Has some konwledge of how certain streams are organised in Mars.
-
+Has some knowledge of how certain streams are organised in Mars.
 """
 
 import datetime
 import logging
 import os
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Tuple
 
 import yaml
 
@@ -30,16 +33,18 @@ DEFAULT_MARS_LABELLING = {
 }
 
 
-def _expand_mars_labelling(request):
+def _expand_mars_labelling(request: Dict[str, Any]) -> Dict[str, Any]:
     """Expand the request with the default Mars labelling.
 
-    The default Mars labelling is:
-
-    {'class': 'od',
-     'type': 'an',
-     'stream': 'oper',
-     'expver': '0001'}
+    Parameters
+    ----------
+    request : dict
+        The original Mars request.
 
+    Returns
+    -------
+    dict
+        The Mars request expanded with default labelling.
     """
     result = DEFAULT_MARS_LABELLING.copy()
     result.update(request)
@@ -49,7 +54,19 @@ def _expand_mars_labelling(request):
 STREAMS = None
 
 
-def _lookup_mars_stream(request):
+def _lookup_mars_stream(request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """Look up the Mars stream information for a given request.
+
+    Parameters
+    ----------
+    request : dict
+        The Mars request.
+
+    Returns
+    -------
+    dict or None
+        The stream information if a match is found, otherwise None.
+    """
     global STREAMS
 
     if STREAMS is None:
@@ -64,8 +81,25 @@ def _lookup_mars_stream(request):
             return s["info"]
 
 
-def recenter(date, center, members):
-
+def recenter(
+    date: datetime.datetime, center: Dict[str, Any], members: Dict[str, Any]
+) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
+    """Recenter the given date with the specified center and members.
+
+    Parameters
+    ----------
+    date : datetime.datetime
+        The date to recenter.
+    center : dict
+        The center request information.
+    members : dict
+        The members request information.
+
+    Returns
+    -------
+    tuple
+        A tuple containing the recentered center and members information.
+    """
     center = _lookup_mars_stream(center)
     members = _lookup_mars_stream(members)
 

anemoi/utils/mars/requests.py

@@ -6,9 +6,23 @@
 # nor does it submit to any jurisdiction.
 
 import sys
+from typing import Any
+from typing import Dict
+from typing import TextIO
 
 
-def print_request(verb, request, file=sys.stdout):
+def print_request(verb: str, request: Dict[str, Any], file: TextIO = sys.stdout) -> None:
+    """Prints a formatted request.
+
+    Parameters
+    ----------
+    verb : str
+        A mars verb
+    request : Dict[str, Any]
+        The request parameters.
+    file : TextIO, optional
+        The file to which the request is printed, by default sys.stdout.
+    """
     r = [verb]
     for k, v in request.items():
         if not isinstance(v, (list, tuple, set)):

anemoi/utils/provenance.py

@@ -14,7 +14,6 @@
  - The versions of the modules which are currently loaded
  - The git information for the modules which are currently loaded from a git repository
  - ...
-
 """
 
 import datetime
@@ -25,11 +24,29 @@ import subprocess
 import sys
 import sysconfig
 from functools import cache
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
 
 LOG = logging.getLogger(__name__)
 
 
-def lookup_git_repo(path):
+def lookup_git_repo(path: str) -> Optional[Any]:
+    """Lookup the git repository for a given path.
+
+    Parameters
+    ----------
+    path : str
+        The path to lookup.
+
+    Returns
+    -------
+    Repo, optional
+        The git repository if found, otherwise None.
+    """
     from git import InvalidGitRepositoryError
     from git import Repo
 
@@ -42,7 +59,21 @@ def lookup_git_repo(path):
     return None
 
 
-def _check_for_git(paths, full):
+def _check_for_git(paths: List[Tuple[str, str]], full: bool) -> Dict[str, Any]:
+    """Check for git information for the given paths.
+
+    Parameters
+    ----------
+    paths : list of tuple
+        The list of paths to check.
+    full : bool
+        Whether to collect full information.
+
+    Returns
+    -------
+    dict
+        The git information for the given paths.
+    """
     versions = {}
     for name, path in paths:
         repo = lookup_git_repo(path)
@@ -77,7 +108,28 @@ def _check_for_git(paths, full):
     return versions
 
 
-def version(versions, name, module, roots, namespaces, paths, full):
+def version(
+    versions: Dict[str, Any], name: str, module: Any, roots: Dict[str, str], namespaces: set, paths: set, full: bool
+) -> None:
+    """Collect version information for a module.
+
+    Parameters
+    ----------
+    versions : dict
+        The dictionary to store the version information.
+    name : str
+        The name of the module.
+    module : Any
+        The module to collect information for.
+    roots : dict
+        The dictionary of root paths.
+    namespaces : set
+        The set of namespaces.
+    paths : set
+        The set of paths.
+    full : bool
+        Whether to collect full information.
+    """
     path = None
 
     if hasattr(module, "__file__"):
@@ -119,7 +171,19 @@ def version(versions, name, module, roots, namespaces, paths, full):
     versions[name] = str(module)
 
 
-def _module_versions(full):
+def _module_versions(full: bool) -> Tuple[Dict[str, Any], set]:
+    """Collect version information for all loaded modules.
+
+    Parameters
+    ----------
+    full : bool
+        Whether to collect full information.
+
+    Returns
+    -------
+    tuple of dict and set
+        The version information and the set of paths.
+    """
     # https://docs.python.org/3/library/sysconfig.html
 
     roots = {}
@@ -149,7 +213,14 @@ def _module_versions(full):
 
 
 @cache
-def package_distributions() -> dict[str, list[str]]:
+def package_distributions() -> Dict[str, List[str]]:
+    """Get the package distributions.
+
+    Returns
+    -------
+    dict
+        The package distributions.
+    """
     # Takes a significant amount of time to run
     # so cache the result
     from importlib import metadata
@@ -161,7 +232,20 @@ def package_distributions() -> dict[str, list[str]]:
     return metadata.packages_distributions()
 
 
-def import_name_to_distribution_name(packages: list):
+def import_name_to_distribution_name(packages: List[str]) -> Dict[str, str]:
+    """Convert import names to distribution names.
+
+    Parameters
+    ----------
+    packages : list of str
+        The list of import names.
+
+    Returns
+    -------
+    dict
+        The dictionary mapping import names to distribution names.
+    """
+
     distribution_names = {}
     package_distribution_names = package_distributions()
 
@@ -179,13 +263,37 @@ def import_name_to_distribution_name(packages: list):
     return distribution_names
 
 
-def module_versions(full):
+def module_versions(full: bool) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Collect version information for all loaded modules and their git information.
+
+    Parameters
+    ----------
+    full : bool
+        Whether to collect full information.
+
+    Returns
+    -------
+    tuple of dict and dict
+        The version information and the git information.
+    """
     versions, paths = _module_versions(full)
     git_versions = _check_for_git(paths, full)
     return versions, git_versions
 
 
-def _name(obj):
+def _name(obj: Any) -> str:
+    """Get the name of an object.
+
+    Parameters
+    ----------
+    obj : Any
+        The object to get the name of.
+
+    Returns
+    -------
+    str
+        The name of the object.
+    """
     if hasattr(obj, "__name__"):
         if hasattr(obj, "__module__"):
             return f"{obj.__module__}.{obj.__name__}"
@@ -195,8 +303,19 @@ def _name(obj):
     return str(obj)
 
 
-def _paths(path_or_object):
+def _paths(path_or_object: Union[None, str, List[str], Tuple[str], Any]) -> List[Tuple[str, str]]:
+    """Get the paths for a given path or object.
 
+    Parameters
+    ----------
+    path_or_object : str, list, tuple, or object
+        The path or object to get the paths for.
+
+    Returns
+    -------
+    list of tuple
+        The list of paths.
+    """
     if path_or_object is None:
         _, paths = _module_versions(full=False)
         return paths
@@ -235,7 +354,7 @@ def _paths(path_or_object):
     return paths
 
 
-def git_check(*args) -> dict:
+def git_check(*args: Any) -> Dict[str, Any]:
     """Return the git information for the given arguments.
 
     Arguments can be:
@@ -255,18 +374,18 @@ def git_check(*args) -> dict:
     dict
         An object with the git information for the given arguments.
 
-    >>> {
-            "anemoi.utils": {
-                "sha1": "c999d83ae283bcbb99f68d92c42d24315922129f",
-                "remotes": [
-                    "git@github.com:ecmwf/anemoi-utils.git"
-                ],
-                "modified_files": [
-                    "anemoi/utils/checkpoints.py"
-                ],
-                "untracked_files": []
+        >>> {
+                "anemoi.utils": {
+                    "sha1": "c999d83ae283bcbb99f68d92c42d24315922129f",
+                    "remotes": [
+                        "git@github.com:ecmwf/anemoi-utils.git"
+                    ],
+                    "modified_files": [
+                        "anemoi/utils/checkpoints.py"
+                    ],
+                    "untracked_files": []
+                }
             }
-        }
     """
     paths = _paths(args if len(args) > 0 else None)
 
@@ -278,7 +397,14 @@ def git_check(*args) -> dict:
     return result
 
 
-def platform_info():
+def platform_info() -> Dict[str, Any]:
+    """Get the platform information.
+
+    Returns
+    -------
+    dict
+        The platform information.
+    """
     import platform
 
     r = {}
@@ -300,7 +426,14 @@ def platform_info():
     return r
 
 
-def gpu_info():
+def gpu_info() -> Union[str, List[Dict[str, Any]]]:
+    """Get the GPU information.
+
+    Returns
+    -------
+    str or list of dict
+        The GPU information or an error message.
+    """
     import nvsmi
 
     if not nvsmi.is_nvidia_smi_on_path():
@@ -312,7 +445,19 @@ def gpu_info():
         return e.output.decode("utf-8").strip()
 
 
-def path_md5(path):
+def path_md5(path: str) -> str:
+    """Calculate the MD5 hash of a file.
+
+    Parameters
+    ----------
+    path : str
+        The path to the file.
+
+    Returns
+    -------
+    str
+        The MD5 hash of the file.
+    """
     import hashlib
 
     hash = hashlib.md5()
@@ -322,7 +467,19 @@ def path_md5(path):
     return hash.hexdigest()
 
 
-def assets_info(paths):
+def assets_info(paths: List[str]) -> Dict[str, Any]:
+    """Get information about the given assets.
+
+    Parameters
+    ----------
+    paths : list of str
+        The list of paths to the assets.
+
+    Returns
+    -------
+    dict
+        The information about the assets.
+    """
     result = {}
 
     for path in paths:
@@ -351,8 +508,8 @@ def assets_info(paths):
     return result
 
 
-def gather_provenance_info(assets=[], full=False) -> dict:
-    """Gather information about the current environment
+def gather_provenance_info(assets: List[str] = [], full: bool = False) -> Dict[str, Any]:
+    """Gather information about the current environment.
 
     Parameters
     ----------

anemoi/utils/registry.py

@@ -12,6 +12,11 @@ import importlib
 import logging
 import os
 import sys
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Optional
+from typing import Union
 
 import entrypoints
 
@@ -19,13 +24,33 @@ LOG = logging.getLogger(__name__)
 
 
 class Wrapper:
-    """A wrapper for the registry"""
+    """A wrapper for the registry.
 
-    def __init__(self, name, registry):
+    Parameters
+    ----------
+    name : str
+        The name of the wrapper.
+    registry : Registry
+        The registry to wrap.
+    """
+
+    def __init__(self, name: str, registry: "Registry"):
         self.name = name
         self.registry = registry
 
-    def __call__(self, factory):
+    def __call__(self, factory: Callable) -> Callable:
+        """Register a factory with the registry.
+
+        Parameters
+        ----------
+        factory : Callable
+            The factory to register.
+
+        Returns
+        -------
+        Callable
+            The registered factory.
+        """
         self.registry.register(self.name, factory)
         return factory
 
@@ -34,10 +59,17 @@ _BY_KIND = {}
 
 
 class Registry:
-    """A registry of factories"""
+    """A registry of factories.
 
-    def __init__(self, package, key="_type"):
+    Parameters
+    ----------
+    package : str
+        The package name.
+    key : str, optional
+        The key to use for the registry, by default "_type".
+    """
 
+    def __init__(self, package: str, key: str = "_type"):
         self.package = package
         self.registered = {}
         self.kind = package.split(".")[-1]
@@ -45,11 +77,36 @@ class Registry:
         _BY_KIND[self.kind] = self
 
     @classmethod
-    def lookup_kind(cls, kind: str):
+    def lookup_kind(cls, kind: str) -> Optional["Registry"]:
+        """Lookup a registry by kind.
+
+        Parameters
+        ----------
+        kind : str
+            The kind of the registry.
+
+        Returns
+        -------
+        Registry, optional
+            The registry if found, otherwise None.
+        """
         return _BY_KIND.get(kind)
 
-    def register(self, name: str, factory: callable = None):
-
+    def register(self, name: str, factory: Optional[Callable] = None) -> Optional[Wrapper]:
+        """Register a factory with the registry.
+
+        Parameters
+        ----------
+        name : str
+            The name of the factory.
+        factory : Callable, optional
+            The factory to register, by default None.
+
+        Returns
+        -------
+        Wrapper, optional
+            A wrapper if the factory is None, otherwise None.
+        """
         if factory is None:
             return Wrapper(name, self)
 
@@ -58,15 +115,35 @@ class Registry:
     # def registered(self, name: str):
     #     return name in self.registered
 
-    def _load(self, file):
+    def _load(self, file: str) -> None:
+        """Load a module from a file.
+
+        Parameters
+        ----------
+        file : str
+            The file to load.
+        """
         name, _ = os.path.splitext(file)
         try:
             importlib.import_module(f".{name}", package=self.package)
         except Exception:
             LOG.warning(f"Error loading filter '{self.package}.{name}'", exc_info=True)
 
-    def lookup(self, name: str, *, return_none=False) -> callable:
-
+    def lookup(self, name: str, *, return_none: bool = False) -> Optional[Callable]:
+        """Lookup a factory by name.
+
+        Parameters
+        ----------
+        name : str
+            The name of the factory.
+        return_none : bool, optional
+            Whether to return None if the factory is not found, by default False.
+
+        Returns
+        -------
+        Callable, optional
+            The factory if found, otherwise None.
+        """
         # print('✅✅✅✅✅✅✅✅✅✅✅✅✅', name, self.registered)
         if name in self.registered:
             return self.registered[name]
@@ -110,14 +187,46 @@ class Registry:
 
         return self.registered[name]
 
-    def create(self, name: str, *args, **kwargs):
+    def create(self, name: str, *args: Any, **kwargs: Any) -> Any:
+        """Create an instance using a factory.
+
+        Parameters
+        ----------
+        name : str
+            The name of the factory.
+        *args : Any
+            Positional arguments for the factory.
+        **kwargs : Any
+            Keyword arguments for the factory.
+
+        Returns
+        -------
+        Any
+            The created instance.
+        """
         factory = self.lookup(name)
         return factory(*args, **kwargs)
 
     # def __call__(self, name: str, *args, **kwargs):
     #     return self.create(name, *args, **kwargs)
 
-    def from_config(self, config, *args, **kwargs):
+    def from_config(self, config: Union[str, Dict[str, Any]], *args: Any, **kwargs: Any) -> Any:
+        """Create an instance from a configuration.
+
+        Parameters
+        ----------
+        config : str or dict
+            The configuration.
+        *args : Any
+            Positional arguments for the factory.
+        **kwargs : Any
+            Keyword arguments for the factory.
+
+        Returns
+        -------
+        Any
+            The created instance.
+        """
         if isinstance(config, str):
             config = {config: {}}