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

anemoi/utils/compatibility.py

@@ -12,10 +12,12 @@ from __future__ import annotations
 import functools
 from typing import Any
 from typing import Callable
+from typing import Optional
+from typing import Union
 
 
 def aliases(
-    aliases: dict[str, str | list[str]] | None = None, **kwargs: str | list[str]
+    aliases: Optional[dict[str, Union[str, list[str]]]] = None, **kwargs: Any
 ) -> Callable[[Callable], Callable]:
     """Alias keyword arguments in a function call.
 
@@ -23,10 +25,10 @@ def aliases(
 
     Parameters
     ----------
-    aliases : dict[str, str  |  list[str]] | None, optional
+    aliases : dict[str, Union[str, list[str]]], optional
         Key, value pair of aliases, with keys being the true name, and value being a str or list of aliases,
         by default None
-    **kwargs : str | list[str]
+    **kwargs : Any
         Kwargs form of aliases
 
     Returns
@@ -49,7 +51,6 @@ def aliases(
     func(a=1, c=2)  # (1, 2)
     func(b=1, d=2)  # (1, 2)
     ```
-
     """
 
     if aliases is None:
@@ -60,7 +61,7 @@ def aliases(
 
     def decorator(func: Callable) -> Callable:
         @functools.wraps(func)
-        def wrapper(*args, **kwargs) -> Any:
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
             keys = kwargs.keys()
             for k in set(keys).intersection(set(aliases.keys())):
                 if aliases[k] in keys:

anemoi/utils/config.py

@@ -14,6 +14,9 @@ import json
 import logging
 import os
 import threading
+from typing import Any
+from typing import Optional
+from typing import Union
 
 import yaml
 
@@ -43,10 +46,18 @@ class DotDict(dict):
     The DotDict class has the same constructor as the dict class.
 
     >>> d = DotDict(a=1, b=2)
-
     """
 
     def __init__(self, *args, **kwargs):
+        """Initialize a DotDict instance.
+
+        Parameters
+        ----------
+        *args : tuple
+            Positional arguments for the dict constructor.
+        **kwargs : dict
+            Keyword arguments for the dict constructor.
+        """
         super().__init__(*args, **kwargs)
 
         for k, v in self.items():
@@ -60,7 +71,19 @@ class DotDict(dict):
                 self[k] = [DotDict(i) if isinstance(i, dict) or is_omegaconf_dict(i) else i for i in v]
 
     @classmethod
-    def from_file(cls, path: str):
+    def from_file(cls, path: str) -> DotDict:
+        """Create a DotDict from a file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the file.
+
+        Returns
+        -------
+        DotDict
+            The created DotDict.
+        """
         _, ext = os.path.splitext(path)
         if ext == ".yaml" or ext == ".yml":
             return cls.from_yaml_file(path)
@@ -72,41 +95,117 @@ class DotDict(dict):
             raise ValueError(f"Unknown file extension {ext}")
 
     @classmethod
-    def from_yaml_file(cls, path: str):
+    def from_yaml_file(cls, path: str) -> DotDict:
+        """Create a DotDict from a YAML file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the YAML file.
+
+        Returns
+        -------
+        DotDict
+            The created DotDict.
+        """
         with open(path, "r") as file:
             data = yaml.safe_load(file)
 
         return cls(data)
 
     @classmethod
-    def from_json_file(cls, path: str):
+    def from_json_file(cls, path: str) -> DotDict:
+        """Create a DotDict from a JSON file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the JSON file.
+
+        Returns
+        -------
+        DotDict
+            The created DotDict.
+        """
         with open(path, "r") as file:
             data = json.load(file)
 
         return cls(data)
 
     @classmethod
-    def from_toml_file(cls, path: str):
+    def from_toml_file(cls, path: str) -> DotDict:
+        """Create a DotDict from a TOML file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the TOML file.
+
+        Returns
+        -------
+        DotDict
+            The created DotDict.
+        """
         with open(path, "r") as file:
             data = tomllib.load(file)
         return cls(data)
 
-    def __getattr__(self, attr):
+    def __getattr__(self, attr: str) -> Any:
+        """Get an attribute.
+
+        Parameters
+        ----------
+        attr : str
+            The attribute name.
+
+        Returns
+        -------
+        Any
+            The attribute value.
+        """
         try:
             return self[attr]
         except KeyError:
             raise AttributeError(attr)
 
-    def __setattr__(self, attr, value):
+    def __setattr__(self, attr: str, value: Any) -> None:
+        """Set an attribute.
+
+        Parameters
+        ----------
+        attr : str
+            The attribute name.
+        value : Any
+            The attribute value.
+        """
         if isinstance(value, dict):
             value = DotDict(value)
         self[attr] = value
 
     def __repr__(self) -> str:
+        """Return a string representation of the DotDict.
+
+        Returns
+        -------
+        str
+            The string representation.
+        """
         return f"DotDict({super().__repr__()})"
 
 
-def is_omegaconf_dict(value) -> bool:
+def is_omegaconf_dict(value: Any) -> bool:
+    """Check if a value is an OmegaConf DictConfig.
+
+    Parameters
+    ----------
+    value : Any
+        The value to check.
+
+    Returns
+    -------
+    bool
+        True if the value is a DictConfig, False otherwise.
+    """
     try:
         from omegaconf import DictConfig
 
@@ -115,7 +214,19 @@ def is_omegaconf_dict(value) -> bool:
         return False
 
 
-def is_omegaconf_list(value) -> bool:
+def is_omegaconf_list(value: Any) -> bool:
+    """Check if a value is an OmegaConf ListConfig.
+
+    Parameters
+    ----------
+    value : Any
+        The value to check.
+
+    Returns
+    -------
+    bool
+        True if the value is a ListConfig, False otherwise.
+    """
     try:
         from omegaconf import ListConfig
 
@@ -130,7 +241,23 @@ CONFIG_LOCK = threading.RLock()
 QUIET = False
 
 
-def _find(config, what, result=None):
+def _find(config: Union[dict, list], what: str, result: list = None) -> list:
+    """Find all occurrences of a key in a nested dictionary or list.
+
+    Parameters
+    ----------
+    config : dict or list
+        The configuration to search.
+    what : str
+        The key to search for.
+    result : list, optional
+        The list to store results, by default None.
+
+    Returns
+    -------
+    list
+        The list of found values.
+    """
     if result is None:
         result = []
 
@@ -149,7 +276,16 @@ def _find(config, what, result=None):
     return result
 
 
-def _merge_dicts(a, b):
+def _merge_dicts(a: dict, b: dict) -> None:
+    """Merge two dictionaries recursively.
+
+    Parameters
+    ----------
+    a : dict
+        The first dictionary.
+    b : dict
+        The second dictionary.
+    """
     for k, v in b.items():
         if k in a and isinstance(a[k], dict) and isinstance(v, dict):
             _merge_dicts(a[k], v)
@@ -157,7 +293,16 @@ def _merge_dicts(a, b):
             a[k] = v
 
 
-def _set_defaults(a, b):
+def _set_defaults(a: dict, b: dict) -> None:
+    """Set default values in a dictionary.
+
+    Parameters
+    ----------
+    a : dict
+        The dictionary to set defaults in.
+    b : dict
+        The dictionary with default values.
+    """
     for k, v in b.items():
         if k in a and isinstance(a[k], dict) and isinstance(v, dict):
             _set_defaults(a[k], v)
@@ -165,7 +310,19 @@ def _set_defaults(a, b):
             a.setdefault(k, v)
 
 
-def config_path(name="settings.toml"):
+def config_path(name: str = "settings.toml") -> str:
+    """Get the path to a configuration file.
+
+    Parameters
+    ----------
+    name : str, optional
+        The name of the configuration file, by default "settings.toml".
+
+    Returns
+    -------
+    str
+        The path to the configuration file.
+    """
     global QUIET
 
     if name.startswith("/") or name.startswith("."):
@@ -197,7 +354,7 @@ def config_path(name="settings.toml"):
     return full
 
 
-def load_any_dict_format(path) -> dict:
+def load_any_dict_format(path: str) -> dict:
     """Load a configuration file in any supported format: JSON, YAML and TOML.
 
     Parameters
@@ -247,8 +404,27 @@ def load_any_dict_format(path) -> dict:
     return open(path).read()
 
 
-def _load_config(name="settings.toml", secrets=None, defaults=None):
+def _load_config(
+    name: str = "settings.toml",
+    secrets: Optional[Union[str, list[str]]] = None,
+    defaults: Optional[Union[str, dict]] = None,
+) -> DotDict:
+    """Load a configuration file.
 
+    Parameters
+    ----------
+    name : str, optional
+        The name of the configuration file, by default "settings.toml".
+    secrets : str or list, optional
+        The name of the secrets file, by default None.
+    defaults : str or dict, optional
+        The name of the defaults file, by default None.
+
+    Returns
+    -------
+    DotDict
+        The loaded configuration.
+    """
     key = json.dumps((name, secrets, defaults), sort_keys=True, default=str)
     if key in CONFIG:
         return CONFIG[key]
@@ -287,7 +463,16 @@ def _load_config(name="settings.toml", secrets=None, defaults=None):
     return CONFIG[key]
 
 
-def _save_config(name, data) -> None:
+def _save_config(name: str, data: Any) -> None:
+    """Save a configuration file.
+
+    Parameters
+    ----------
+    name : str
+        The name of the configuration file.
+    data : Any
+        The data to save.
+    """
     CONFIG.pop(name, None)
 
     conf = config_path(name)
@@ -309,7 +494,7 @@ def _save_config(name, data) -> None:
         f.write(data)
 
 
-def save_config(name, data) -> None:
+def save_config(name: str, data: Any) -> None:
     """Save a configuration file.
 
     Parameters
@@ -319,13 +504,16 @@ def save_config(name, data) -> None:
 
     data : Any
         The data to save.
-
     """
     with CONFIG_LOCK:
         _save_config(name, data)
 
 
-def load_config(name="settings.toml", secrets=None, defaults=None) -> DotDict | str:
+def load_config(
+    name: str = "settings.toml",
+    secrets: Optional[Union[str, list[str]]] = None,
+    defaults: Optional[Union[str, dict]] = None,
+) -> DotDict | str:
     """Read a configuration file.
 
     Parameters
@@ -347,8 +535,21 @@ def load_config(name="settings.toml", secrets=None, defaults=None) -> DotDict |
         return _load_config(name, secrets, defaults)
 
 
-def load_raw_config(name, default=None) -> DotDict | str:
+def load_raw_config(name: str, default: Any = None) -> Union[DotDict, str]:
+    """Load a raw configuration file.
+
+    Parameters
+    ----------
+    name : str
+        The name of the configuration file.
+    default : Any, optional
+        The default value if the file does not exist, by default None.
 
+    Returns
+    -------
+    DotDict or str
+        The loaded configuration or the default value.
+    """
     path = config_path(name)
     if os.path.exists(path):
         return load_any_dict_format(path)
@@ -356,7 +557,7 @@ def load_raw_config(name, default=None) -> DotDict | str:
     return default
 
 
-def check_config_mode(name="settings.toml", secrets_name=None, secrets=None) -> None:
+def check_config_mode(name: str = "settings.toml", secrets_name: str = None, secrets: list[str] = None) -> None:
     """Check that a configuration file is secure.
 
     Parameters
@@ -393,7 +594,25 @@ def check_config_mode(name="settings.toml", secrets_name=None, secrets=None) ->
         CHECKED[name] = True
 
 
-def find(metadata, what, result=None, *, select: callable = None):
+def find(metadata: Union[dict, list], what: str, result: list = None, *, select: callable = None) -> list:
+    """Find all occurrences of a key in a nested dictionary or list with an optional selector.
+
+    Parameters
+    ----------
+    metadata : dict or list
+        The metadata to search.
+    what : str
+        The key to search for.
+    result : list, optional
+        The list to store results, by default None.
+    select : callable, optional
+        A function to filter the results, by default None.
+
+    Returns
+    -------
+    list
+        The list of found values.
+    """
     if result is None:
         result = []
 
@@ -413,7 +632,19 @@ def find(metadata, what, result=None, *, select: callable = None):
     return result
 
 
-def merge_configs(*configs):
+def merge_configs(*configs: dict) -> dict:
+    """Merge multiple configuration dictionaries.
+
+    Parameters
+    ----------
+    *configs : dict
+        The configuration dictionaries to merge.
+
+    Returns
+    -------
+    dict
+        The merged configuration dictionary.
+    """
     result = {}
     for config in configs:
         _merge_dicts(result, config)