anemoi-utils 0.1.2__tar.gz → 0.1.3__tar.gz

{anemoi-utils-0.1.2 → anemoi-utils-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.2
+Version: 0.1.3
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Home-page: https://github.com/ecmwf/anemoi-utils
 Author: European Centre for Medium-Range Weather Forecasts (ECMWF)

{anemoi-utils-0.1.2 → anemoi-utils-0.1.3}/anemoi/utils/__init__.py

@@ -6,4 +6,4 @@
 # nor does it submit to any jurisdiction.
 
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"

{anemoi-utils-0.1.2 → anemoi-utils-0.1.3}/anemoi/utils/checkpoints.py

@@ -6,6 +6,11 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+"""
+Read and write extra metadata in PyTorch checkpoints files. These files
+are zip archives containing the model weights.
+"""
+
 import json
 import logging
 import os
@@ -16,7 +21,26 @@ LOG = logging.getLogger(__name__)
 DEFAULT_NAME = "anemoi-metadata.json"
 
 
-def load_metadata(path, name=DEFAULT_NAME):
+def load_metadata(path: str, name: str = DEFAULT_NAME):
+    """Load 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
+
+    Returns
+    -------
+    JSON
+        The content of the metadata file
+
+    Raises
+    ------
+    ValueError
+        If the metadata file is not found
+    """
     with zipfile.ZipFile(path, "r") as f:
         metadata = None
         for b in f.namelist():
@@ -33,6 +57,17 @@ def load_metadata(path, name=DEFAULT_NAME):
 
 
 def save_metadata(path, metadata, name=DEFAULT_NAME):
+    """Save metadata to a checkpoint file
+
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    metadata : JSON
+        A JSON serializable object
+    name : str, optional
+        The name of the metadata file in the zip archive
+    """
     with zipfile.ZipFile(path, "a") as zipf:
         base, _ = os.path.splitext(os.path.basename(path))
         zipf.writestr(

{anemoi-utils-0.1.2 → anemoi-utils-0.1.3}/anemoi/utils/humanize.py

@@ -7,18 +7,41 @@
 # nor does it submit to any jurisdiction.
 #
 
+"""
+Generate human readable strings
+"""
+
 import datetime
 import re
 from collections import defaultdict
 
 
-def bytes(n):
-    """
+def bytes(n: float) -> str:
+    """Convert a number of bytes to a human readable string
+
     >>> bytes(4096)
     '4 KiB'
+
     >>> bytes(4000)
     '3.9 KiB'
+
+    Parameters
+    ----------
+    n : float
+        the number of bytes
+
+    Returns
+    -------
+    str
+        a human readable string
+    """
+
     """
+
+
+
+    """
+
     if n < 0:
         sign = "-"
         n -= 0
@@ -33,13 +56,7 @@ def bytes(n):
     return "%s%g%s" % (sign, int(n * 10 + 0.5) / 10.0, u[i])
 
 
-def base2(n):
-    """
-    >>> base2(4096)
-    '4K'
-    >>> base2(4000)
-    '3.9K'
-    """
+def base2(n) -> str:
 
     u = ["", "K", "M", "G", "T", " P", "E", "Z", "Y"]
     i = 0
@@ -65,7 +82,24 @@ def _plural(count):
         return ""
 
 
-def seconds(seconds):
+def seconds(seconds: float) -> str:
+    """Convert a number of seconds to a human readable string
+
+    >>> seconds(4000)
+    '1 hour 6 minutes 40 seconds'
+
+
+    Parameters
+    ----------
+    seconds : float
+        The number of seconds
+
+    Returns
+    -------
+    str
+        A human readable string
+
+    """
     if isinstance(seconds, datetime.timedelta):
         seconds = seconds.total_seconds()
 
@@ -113,6 +147,20 @@ def number(value):
 
 
 def plural(value, what):
+    """_summary_
+
+    Parameters
+    ----------
+    value : _type_
+        _description_
+    what : _type_
+        _description_
+
+    Returns
+    -------
+    _type_
+        _description_
+    """
     return f"{number(value)} {what}{_plural(value)}"
 
 
@@ -159,6 +207,39 @@ def __(n):
 
 
 def when(then, now=None, short=True):
+    """Generate a human readable string for a date, relative to now
+
+
+    >>> when(datetime.datetime.now() - datetime.timedelta(hours=2))
+    '2 hours ago'
+
+    >>> when(datetime.datetime.now() - datetime.timedelta(days=1))
+    'yesterday at 08:46'
+
+    >>> when(datetime.datetime.now() - datetime.timedelta(days=5))
+    'last Sunday'
+
+    >>> when(datetime.datetime.now() - datetime.timedelta(days=365))
+    'last year'
+
+    >>> when(datetime.datetime.now() + datetime.timedelta(days=365))
+    'next year'
+
+    Parameters
+    ----------
+    then : datetime.datetime
+        A datetime
+    now : datetime.datetime, optional
+        The reference date, by default NOW
+    short : bool, optional
+        Genererate shorter strings, by default True
+
+    Returns
+    -------
+    str
+        A human readable string
+
+    """
     last = "last"
 
     if now is None:

{anemoi-utils-0.1.2 → anemoi-utils-0.1.3}/anemoi/utils/provenance.py

@@ -5,6 +5,16 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+"""
+Collect information about the current environment, like:
+
+ - The Python version
+ - 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
 import json
 import logging
@@ -29,7 +39,7 @@ def lookup_git_repo(path):
     return None
 
 
-def check_for_git(paths, full):
+def _check_for_git(paths, full):
     versions = {}
     for name, path in paths:
         repo = lookup_git_repo(path)
@@ -106,7 +116,7 @@ def version(versions, name, module, roots, namespaces, paths, full):
     versions[name] = str(module)
 
 
-def module_versions(full):
+def _module_versions(full):
     # https://docs.python.org/3/library/sysconfig.html
 
     roots = {}
@@ -131,11 +141,104 @@ def module_versions(full):
         if len(bits) == 2 and bits[0] in namespaces:
             version(versions, k, v, roots, namespaces, paths, full)
 
-    git_versions = check_for_git(paths, full)
+    return versions, paths
 
+
+def module_versions(full):
+    versions, paths = _module_versions(full)
+    git_versions = _check_for_git(paths, full)
     return versions, git_versions
 
 
+def _name(obj):
+    if hasattr(obj, "__name__"):
+        if hasattr(obj, "__module__"):
+            return f"{obj.__module__}.{obj.__name__}"
+        return obj.__name__
+    if hasattr(obj, "__class__"):
+        return _name(obj.__class__)
+    return str(obj)
+
+
+def _paths(path_or_object):
+
+    if path_or_object is None:
+        _, paths = _module_versions(full=False)
+        return paths
+
+    if isinstance(path_or_object, (list, tuple, set)):
+        paths = []
+        for p in path_or_object:
+            paths.extend(_paths(p))
+        return paths
+
+    if isinstance(path_or_object, str):
+        module = sys.modules.get(path_or_object)
+        if module is not None:
+            return _paths(module)
+        return [(path_or_object, path_or_object)]
+
+    if hasattr(path_or_object, "__module__"):
+        module = sys.modules.get(path_or_object.__module__)
+        return [(path_or_object.__module__, module.__file__)]
+
+    name = _name(path_or_object)
+    paths = []
+    if hasattr(path_or_object, "__file__"):
+        paths.append((name, path_or_object.__file__))
+
+    if hasattr(path_or_object, "__code__"):
+        paths.append((name, path_or_object.__code__.co_filename))
+
+    if hasattr(path_or_object, "__module__"):
+        module = sys.modules.get(path_or_object.__module__)
+        paths.append((name, module.__file__))
+
+    if not paths:
+        raise ValueError(f"Could not find path for {name} {path_or_object} {type(path_or_object)}")
+
+    return paths
+
+
+def git_check(*args):
+    """Return the git information for the given arguments.
+
+    Arguments can be:
+        - an empty list, in that case all loaded modules are checked
+        - a module name
+        - a module object
+        - an object or a class
+        - a path to a directory
+
+    Returns
+    -------
+    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": []
+            }
+        }
+    """
+    paths = _paths(args if len(args) > 0 else None)
+
+    git = _check_for_git(paths, full=True)
+    result = {}
+    for k, v in git.items():
+        result[k] = v["git"]
+
+    return result
+
+
 def platform_info():
     import platform
 
@@ -209,7 +312,21 @@ def assets_info(paths):
     return result
 
 
-def gather_provenance_info(assets=[], full=False):
+def gather_provenance_info(assets=[], full=False) -> dict:
+    """Gather information about the current environment
+
+    Parameters
+    ----------
+    assets : list, optional
+        A list of file paths for which to collect the MD5 sum, the size and time attributes, by default []
+    full : bool, optional
+        If true, will also collect various paths, by default False
+
+    Returns
+    -------
+    dict
+        A dictionary with the collected information
+    """
     executable = sys.executable
 
     versions, git_versions = module_versions(full)

{anemoi-utils-0.1.2 → anemoi-utils-0.1.3}/anemoi/utils/text.py

@@ -5,20 +5,60 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+"""
+Text utilities
+"""
 
 import sys
 from collections import defaultdict
 
-from termcolor import colored
-
 # https://en.wikipedia.org/wiki/Box-drawing_character
 
 
-def dotted_line(n=84, file=sys.stdout):
-    print("┈" * n, file=file)
+def dotted_line(n=84) -> str:
+    """_summary_
+
+    Parameters
+    ----------
+    n : int, optional
+        _description_, by default 84
+
+    Returns
+    -------
+    str
+        _description_
+    """
+    return "┈" * n
+
+
+def boxed(text, min_width=80, max_width=80) -> str:
+    """Put a box around a text
+
+    >>> print(boxed("Hello,\\nWorld!", max_width=40))
+    ┌──────────────────────────────────────────┐
+    │ Hello,                                   │
+    │ World!                                   │
+    └──────────────────────────────────────────┘
+
+    Parameters
+    ----------
+    text : str
+        The text to box
+    min_width : int, optional
+        The minimum width of the box, by default 80
+    max_width : int, optional
+        The maximum width of the box, by default 80
+
+    Returns
+    -------
+    str
+        A boxed version of the input text
+
+
 
 
-def boxed(text, min_width=80, max_width=80):
+    """
+
     lines = text.split("\n")
     width = max(len(_) for _ in lines)
 
@@ -44,14 +84,20 @@ def boxed(text, min_width=80, max_width=80):
 
 
 def bold(text):
+    from termcolor import colored
+
     return colored(text, attrs=["bold"])
 
 
 def red(text):
+    from termcolor import colored
+
     return colored(text, "red")
 
 
 def green(text):
+    from termcolor import colored
+
     return colored(text, "green")
 
 
@@ -191,6 +237,20 @@ class Tree:
 
 
 def table(rows, header, align, margin=0):
+    """_summary_
+
+    Parameters
+    ----------
+    rows : _type_
+        _description_
+    header : _type_
+        _description_
+    align : _type_
+        _description_
+    margin : int, optional
+        _description_, by default 0
+    """
+
     def _(x):
         try:
             x = float(x)
@@ -231,6 +291,27 @@ def table(rows, header, align, margin=0):
     return "\n".join(result)
 
 
-def progress(done, todo, width=80):
+def progress(done, todo, width=80) -> str:
+    """_summary_
+
+    >>> print(progress(10, 100,width=50))
+    █████▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒
+
+    Parameters
+    ----------
+    done : function
+        _description_
+    todo : _type_
+        _description_
+    width : int, optional
+        _description_, by default 80
+
+    Returns
+    -------
+    str
+        _description_
+
+
+    """
     done = min(int(done / todo * width + 0.5), width)
     return green("█" * done) + red("█" * (width - done))

{anemoi-utils-0.1.2 → anemoi-utils-0.1.3}/anemoi_utils.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.2
+Version: 0.1.3
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Home-page: https://github.com/ecmwf/anemoi-utils
 Author: European Centre for Medium-Range Weather Forecasts (ECMWF)

{anemoi-utils-0.1.2 → anemoi-utils-0.1.3}/setup.py

@@ -41,6 +41,7 @@ provenance_requires = [
 text_requires = [
     "termcolor",
 ]
+
 doc_requires = ["sphinx", "sphinx_rtd_theme", "nbsphinx", "pandoc"]
 
 all_requires = install_requires + provenance_requires + text_requires