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

anemoi/utils/s3.py

@@ -8,6 +8,9 @@
 # nor does it submit to any jurisdiction.
 
 import warnings
+from typing import Any
+from typing import Callable
+from typing import Optional
 
 from .remote import transfer
 from .remote.s3 import delete as delete_
@@ -21,7 +24,21 @@ warnings.warn(
 )
 
 
-def s3_client(*args, **kwargs):
+def s3_client(*args: Any, **kwargs: Any) -> Any:
+    """Create an S3 client.
+
+    Parameters
+    ----------
+    *args : Any
+        Positional arguments for the S3 client.
+    **kwargs : Any
+        Keyword arguments for the S3 client.
+
+    Returns
+    -------
+    Any
+        The S3 client.
+    """
     warnings.warn(
         "The 's3_client' function (from anemoi.utils.s3 import s3_client) function is deprecated and will be removed in a future release. "
         "Please use the 's3_client' function (from anemoi.utils.remote.s3 import s3_client) instead.",
@@ -31,7 +48,35 @@ def s3_client(*args, **kwargs):
     return s3_client_(*args, **kwargs)
 
 
-def upload(source, target, *, overwrite=False, resume=False, verbosity=1, progress=None, threads=1) -> None:
+def upload(
+    source: str,
+    target: str,
+    *,
+    overwrite: bool = False,
+    resume: bool = False,
+    verbosity: int = 1,
+    progress: Optional[Callable] = None,
+    threads: int = 1,
+) -> None:
+    """Upload a file to S3.
+
+    Parameters
+    ----------
+    source : str
+        The source file path.
+    target : str
+        The target S3 path.
+    overwrite : bool, optional
+        Whether to overwrite the target file, by default False.
+    resume : bool, optional
+        Whether to resume a previous upload, by default False.
+    verbosity : int, optional
+        The verbosity level, by default 1.
+    progress : Callable, optional
+        A callback function for progress updates, by default None.
+    threads : int, optional
+        The number of threads to use, by default 1.
+    """
     warnings.warn(
         "The 'upload' function (from anemoi.utils.s3 import upload) function is deprecated and will be removed in a future release. "
         "Please use the 'transfer' function (from anemoi.utils.remote import transfer) instead.",
@@ -43,7 +88,21 @@ def upload(source, target, *, overwrite=False, resume=False, verbosity=1, progre
     )
 
 
-def download(*args, **kwargs):
+def download(*args: Any, **kwargs: Any) -> Any:
+    """Download a file from S3.
+
+    Parameters
+    ----------
+    *args : Any
+        Positional arguments for the download.
+    **kwargs : Any
+        Keyword arguments for the download.
+
+    Returns
+    -------
+    Any
+        The result of the download.
+    """
     warnings.warn(
         "The 'download' function (from anemoi.utils.s3 import download) function is deprecated and will be removed in a future release. "
         "Please use the 'transfer' function (from anemoi.utils.remote import transfer) instead.",
@@ -53,7 +112,21 @@ def download(*args, **kwargs):
     return transfer(*args, **kwargs)
 
 
-def delete(*args, **kwargs):
+def delete(*args: Any, **kwargs: Any) -> Any:
+    """Delete a file from S3.
+
+    Parameters
+    ----------
+    *args : Any
+        Positional arguments for the delete.
+    **kwargs : Any
+        Keyword arguments for the delete.
+
+    Returns
+    -------
+    Any
+        The result of the delete.
+    """
     warnings.warn(
         "The 'delete' function (from anemoi.utils.s3 import delete) function is deprecated and will be removed in a future release. "
         "Please use the 'transfer' function (from anemoi.utils.remote.s3 import delete) instead.",

anemoi/utils/sanitise.py

@@ -11,6 +11,7 @@
 import os
 import re
 from pathlib import Path
+from typing import Any
 from urllib.parse import parse_qs
 from urllib.parse import urlencode
 from urllib.parse import urlparse
@@ -22,10 +23,18 @@ RE1 = re.compile(r"{([^}]*)}")
 RE2 = re.compile(r"\(([^}]*)\)")
 
 
-def sanitise(obj):
-    """sanitise an object:
-    - by replacing all full paths with shortened versions.
-    - by replacing URL passwords with '***'.
+def sanitise(obj: Any) -> Any:
+    """Sanitise an object by replacing all full paths with shortened versions and URL passwords with '***'.
+
+    Parameters
+    ----------
+    obj : Any
+        The object to sanitise.
+
+    Returns
+    -------
+    Any
+        The sanitised object.
     """
 
     if isinstance(obj, dict):
@@ -43,7 +52,19 @@ def sanitise(obj):
     return obj
 
 
-def _sanitise_string(obj):
+def _sanitise_string(obj: str) -> str:
+    """Sanitise a string by replacing full paths and URL passwords.
+
+    Parameters
+    ----------
+    obj : str
+        The string to sanitise.
+
+    Returns
+    -------
+    str
+        The sanitised string.
+    """
 
     parsed = urlparse(obj, allow_fragments=True)
 
@@ -56,7 +77,19 @@ def _sanitise_string(obj):
     return obj
 
 
-def _sanitise_url(parsed):
+def _sanitise_url(parsed: Any) -> str:
+    """Sanitise a URL by replacing passwords with '***'.
+
+    Parameters
+    ----------
+    parsed : Any
+        The parsed URL.
+
+    Returns
+    -------
+    str
+        The sanitised URL.
+    """
 
     LIST = [
         "pass",
@@ -100,7 +133,19 @@ def _sanitise_url(parsed):
     return urlunparse([scheme, netloc, path, params, query, fragment])
 
 
-def _sanitise_path(path):
+def _sanitise_path(path: str) -> str:
+    """Sanitise a file path by shortening it.
+
+    Parameters
+    ----------
+    path : str
+        The file path to sanitise.
+
+    Returns
+    -------
+    str
+        The sanitised file path.
+    """
     bits = list(reversed(Path(path).parts))
     result = [bits.pop(0)]
     for bit in bits:

anemoi/utils/text.py

@@ -8,16 +8,21 @@
 # nor does it submit to any jurisdiction.
 
 
-"""Text utilities"""
+"""Text utilities."""
 
 import re
 from collections import defaultdict
+from typing import Any
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
 
 # https://en.wikipedia.org/wiki/Box-drawing_character
 
 
 def dotted_line(width=84) -> str:
-    """Return a dotted line using '┈'
+    """Return a dotted line using '┈'.
 
     >>> dotted_line(40)
     ┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈
@@ -40,12 +45,35 @@ def dotted_line(width=84) -> str:
 _ansi_escape = re.compile(r"\x1b\[([0-9;]*[mGKH])")
 
 
-def _has_ansi_escape(s):
+def _has_ansi_escape(s: str) -> bool:
+    """Check if a string contains ANSI escape codes.
+
+    Parameters
+    ----------
+    s : str
+        The string to check.
+
+    Returns
+    -------
+    bool
+        True if the string contains ANSI escape codes, False otherwise.
+    """
     return _ansi_escape.search(s) is not None
 
 
-def _split_tokens(s):
-    """Split a string into a list of visual characters with their lenghts."""
+def _split_tokens(s: str) -> List[Tuple[str, int]]:
+    """Split a string into a list of visual characters with their lengths.
+
+    Parameters
+    ----------
+    s : str
+        The string to split.
+
+    Returns
+    -------
+    list of tuple
+        A list of tuples where each tuple contains a visual character and its length.
+    """
     from wcwidth import wcswidth
 
     initial = s
@@ -76,8 +104,19 @@ def _split_tokens(s):
     return out
 
 
-def visual_len(s):
-    """Compute the length of a string as it appears on the terminal."""
+def visual_len(s: Union[str, List[Tuple[str, int]]]) -> int:
+    """Compute the length of a string as it appears on the terminal.
+
+    Parameters
+    ----------
+    s : str or list of tuple
+        The string or list of visual characters with their lengths.
+
+    Returns
+    -------
+    int
+        The visual length of the string.
+    """
     if isinstance(s, str):
         s = _split_tokens(s)
     assert isinstance(s, (tuple, list)), (type(s), s)
@@ -92,8 +131,8 @@ def visual_len(s):
     return n
 
 
-def boxed(text, min_width=80, max_width=80) -> str:
-    """Put a box around a text
+def boxed(text: str, min_width: int = 80, max_width: int = 80) -> str:
+    """Put a box around a text.
 
     >>> boxed("Hello,\\nWorld!", max_width=40)
     ┌──────────────────────────────────────────┐
@@ -114,12 +153,11 @@ def boxed(text, min_width=80, max_width=80) -> str:
     -------
     str
         A boxed version of the input text
-
     """
 
     lines = []
     for line in text.split("\n"):
-        line = line.strip()
+        line = line.rstrip()
         line = _split_tokens(line)
         lines.append(line)
 
@@ -158,68 +196,157 @@ def boxed(text, min_width=80, max_width=80) -> str:
     return "\n".join(box)
 
 
-def bold(text):
+def bold(text: str) -> str:
+    """Make the text bold.
+
+    Parameters
+    ----------
+    text : str
+        The text to make bold.
+
+    Returns
+    -------
+    str
+        The bold text.
+    """
     from termcolor import colored
 
     return colored(text, attrs=["bold"])
 
 
-def red(text):
+def red(text: str) -> str:
+    """Make the text red.
+
+    Parameters
+    ----------
+    text : str
+        The text to make red.
+
+    Returns
+    -------
+    str
+        The red text.
+    """
     from termcolor import colored
 
     return colored(text, "red")
 
 
-def green(text):
+def green(text: str) -> str:
+    """Make the text green.
+
+    Parameters
+    ----------
+    text : str
+        The text to make green.
+
+    Returns
+    -------
+    str
+        The green text.
+    """
     from termcolor import colored
 
     return colored(text, "green")
 
 
-def blue(text):
+def blue(text: str) -> str:
+    """Make the text blue.
+
+    Parameters
+    ----------
+    text : str
+        The text to make blue.
+
+    Returns
+    -------
+    str
+        The blue text.
+    """
     from termcolor import colored
 
     return colored(text, "blue")
 
 
 class Tree:
-    """Tree data structure."""
+    """Tree data structure.
 
-    def __init__(self, actor, parent=None):
+    Parameters
+    ----------
+    actor : Any
+        The actor associated with the tree node.
+    parent : Tree, optional
+        The parent tree node, by default None.
+    """
+
+    def __init__(self, actor: Any, parent: Optional["Tree"] = None):
         self._actor = actor
         self._kids = []
         self._parent = parent
 
-    def adopt(self, kid):
+    def adopt(self, kid: "Tree") -> None:
+        """Adopt a child tree node.
+
+        Parameters
+        ----------
+        kid : Tree
+            The child tree node to adopt.
+        """
         kid._parent._kids.remove(kid)
         self._kids.append(kid)
         kid._parent = self
         # assert False
 
-    def forget(self):
+    def forget(self) -> None:
+        """Forget the current tree node."""
         self._parent._kids.remove(self)
         self._parent = None
 
     @property
-    def is_leaf(self):
+    def is_leaf(self) -> bool:
+        """Bool: True if the tree node is a leaf, False otherwise."""
         return len(self._kids) == 0
 
     @property
-    def key(self):
+    def key(self) -> Tuple:
+        """Tuple: The key of the tree node."""
         return tuple(sorted(self._actor.as_dict().items()))
 
     @property
-    def _text(self):
+    def _text(self) -> str:
+        """Str: The text representation of the tree node."""
         return self._actor.summary
 
     @property
-    def summary(self):
+    def summary(self) -> str:
+        """Str: The summary of the tree node."""
         return self._actor.summary
 
-    def as_dict(self):
+    def as_dict(self) -> dict:
+        """Convert the tree node to a dictionary.
+
+        Returns
+        -------
+        dict
+            The dictionary representation of the tree node.
+        """
         return self._actor.as_dict()
 
-    def node(self, actor, insert=False):
+    def node(self, actor: Any, insert: bool = False) -> "Tree":
+        """Create a new tree node.
+
+        Parameters
+        ----------
+        actor : Any
+            The actor associated with the new tree node.
+        insert : bool, optional
+            Whether to insert the new tree node at the beginning, by default False.
+
+        Returns
+        -------
+        Tree
+            The new tree node.
+        """
         node = Tree(actor, self)
         if insert:
             self._kids.insert(0, node)
@@ -227,7 +354,8 @@ class Tree:
             self._kids.append(node)
         return node
 
-    def print(self):
+    def print(self) -> None:
+        """Print the tree."""
         padding = []
 
         while self._factorise():
@@ -235,14 +363,28 @@ class Tree:
 
         self._print(padding)
 
-    def _leaves(self, result):
+    def _leaves(self, result: List["Tree"]) -> None:
+        """Collect all leaf nodes.
+
+        Parameters
+        ----------
+        result : list of Tree
+            The list to collect the leaf nodes.
+        """
         if self.is_leaf:
             result.append(self)
         else:
             for kid in self._kids:
                 kid._leaves(result)
 
-    def _factorise(self):
+    def _factorise(self) -> bool:
+        """Factorise the tree.
+
+        Returns
+        -------
+        bool
+            True if the tree was factorised, False otherwise.
+        """
         if len(self._kids) == 0:
             return False
 
@@ -290,7 +432,14 @@ class Tree:
 
         return result
 
-    def _print(self, padding):
+    def _print(self, padding: List[str]) -> None:
+        """Print the tree with padding.
+
+        Parameters
+        ----------
+        padding : list of str
+            The padding for each level of the tree.
+        """
         for i, p in enumerate(padding[:-1]):
             if p == " └":
                 padding[i] = "  "
@@ -308,7 +457,19 @@ class Tree:
 
         padding.pop()
 
-    def to_json(self, depth=0):
+    def to_json(self, depth: int = 0) -> dict:
+        """Convert the tree to a JSON serializable dictionary.
+
+        Parameters
+        ----------
+        depth : int, optional
+            The depth of the tree, by default 0.
+
+        Returns
+        -------
+        dict
+            The JSON serializable dictionary representation of the tree.
+        """
         while self._factorise():
             pass
 
@@ -319,20 +480,20 @@ class Tree:
         }
 
 
-def table(rows, header, align, margin=0) -> str:
-    """Format a table
+def table(rows: List[List[Any]], header: List[str], align: List[str], margin: int = 0) -> str:
+    """Format a table.
 
-    >>> table([['Aa', 12, 5],
-               ['B', 120, 1],
-               ['C', 9, 123]],
-               ['C1', 'C2', 'C3'],
-               ['<', '>', '>'])
-        C1 │  C2 │  C3
-        ───┼─────┼────
-        Aa │  12 │   5
-        B  │ 120 │   1
-        C  │   9 │ 123
-        ───┴─────┴────
+        >>> table([['Aa', 12, 5],
+                ['B', 120, 1],
+                ['C', 9, 123]],
+                ['C1', 'C2', 'C3'],
+                ['<', '>', '>'])
+            C1 │  C2 │  C3
+            ───┼─────┼────
+            Aa │  12 │   5
+            B  │ 120 │   1
+            C  │   9 │ 123
+            ───┴─────┴────
 
     Parameters
     ----------
@@ -397,26 +558,29 @@ def table(rows, header, align, margin=0) -> str:
     return "\n".join(result)
 
 
-def progress(done, todo, width=80) -> str:
-    """_summary_
-
-    >>> print(progress(10, 100,width=50))
-    █████▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒
+def progress(done: int, todo: int, width: int = 80) -> str:
+    """Generates a progress bar string.
 
     Parameters
     ----------
-    done : function
-        _description_
-    todo : _type_
-        _description_
+
+    done : int
+        The number of tasks completed.
+    todo : int
+        The total number of tasks.
     width : int, optional
-        _description_, by default 80
+        The width of the progress bar, by default 80.
 
     Returns
     -------
     str
-        _description_
+        A string representing the progress bar.
 
+    Example
+    -------
+
+    >>> print(progress(10, 100,width=50))
+    █████▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒
     """
     done = min(int(done / todo * width + 0.5), width)
     return green("█" * done) + red("█" * (width - done))