anemoi-utils 0.4.12__py3-none-any.whl → 0.4.14__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/testing.py

@@ -0,0 +1,182 @@
+# (C) Copyright 2025- Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import atexit
+import logging
+import os
+import shutil
+import tempfile
+import threading
+
+from multiurl import download
+
+LOG = logging.getLogger(__name__)
+
+TEST_DATA_URL = "https://object-store.os-api.cci1.ecmwf.int/ml-tests/test-data/samples/"
+
+lock = threading.RLock()
+TEMPORARY_DIRECTORY = None
+
+
+def _temporary_directory() -> str:
+    """Return a temporary directory in which to download test data.
+
+    Returns
+    -------
+    str
+        The path to the temporary directory.
+    """
+    global TEMPORARY_DIRECTORY
+    with lock:
+        if TEMPORARY_DIRECTORY is not None:
+            return TEMPORARY_DIRECTORY
+
+        TEMPORARY_DIRECTORY = tempfile.mkdtemp()
+
+        # Register a cleanup function to remove the directory at exit
+        atexit.register(shutil.rmtree, TEMPORARY_DIRECTORY)
+
+        return TEMPORARY_DIRECTORY
+
+
+def _check_path(path: str) -> None:
+    """Check if the given path is normalized, not absolute, and does not start with a dot.
+
+    Parameters
+    ----------
+    path : str
+        The path to check.
+
+    Raises
+    ------
+    AssertionError
+        If the path is not normalized, is absolute, or starts with a dot.
+    """
+    assert os.path.normpath(path) == path, f"Path '{path}' should be normalized"
+    assert not os.path.isabs(path), f"Path '{path}' should not be absolute"
+    assert not path.startswith("."), f"Path '{path}' should not start with '.'"
+
+
+def url_for_test_data(path: str) -> str:
+    """Generate the URL for the test data based on the given path.
+
+    Parameters
+    ----------
+    path : str
+        The relative path to the test data.
+
+    Returns
+    -------
+    str
+        The full URL to the test data.
+    """
+    _check_path(path)
+
+    return f"{TEST_DATA_URL}{path}"
+
+
+def get_test_data(path: str, gzipped=False) -> str:
+    """Download the test data to a temporary directory and return the local path.
+
+    Parameters
+    ----------
+    path : str
+        The relative path to the test data.
+    gzipped : bool, optional
+        Flag indicating if the remote file is gzipped, by default False. The local file will be gunzipped.
+
+    Returns
+    -------
+    str
+        The local path to the downloaded test data.
+    """
+    _check_path(path)
+
+    target = os.path.normpath(os.path.join(_temporary_directory(), path))
+    with lock:
+        if os.path.exists(target):
+            return target
+
+        os.makedirs(os.path.dirname(target), exist_ok=True)
+        url = url_for_test_data(path)
+
+        if gzipped:
+            url += ".gz"
+            target += ".gz"
+
+        LOG.info(f"Downloading test data from {url} to {target}")
+
+        download(url, target)
+
+        if gzipped:
+            import gzip
+
+            with gzip.open(target, "rb") as f_in:
+                with open(target[:-3], "wb") as f_out:
+                    shutil.copyfileobj(f_in, f_out)
+            os.remove(target)
+            target = target[:-3]
+
+        return target
+
+
+def get_test_archive(path: str, extension=".extracted") -> str:
+    """Download an archive file (.zip, .tar, .tar.gz, .tar.bz2, .tar.xz) to a temporary directory
+    unpack it, and return the local path to the directory containing the extracted files.
+
+    Parameters
+    ----------
+    path : str
+        The relative path to the test data.
+    extension : str, optional
+        The extension to add to the extracted directory, by default '.extracted'
+
+    Returns
+    -------
+    str
+        The local path to the downloaded test data.
+    """
+
+    with lock:
+
+        archive = get_test_data(path)
+        target = archive + extension
+
+        shutil.unpack_archive(archive, os.path.dirname(target) + ".tmp")
+        os.rename(os.path.dirname(target) + ".tmp", target)
+
+        return target
+
+
+def packages_installed(*names) -> bool:
+    """Check if all the given packages are installed.
+
+    Use this function to check if the required packages are installed before running tests.
+
+    >>> @pytest.mark.skipif(not packages_installed("foo", "bar"), reason="Packages 'foo' and 'bar' are not installed")
+    >>> def test_foo_bar() -> None:
+    >>>    ...
+
+    Parameters
+    ----------
+    names : str
+        The names of the packages to check.
+
+    Returns
+    -------
+    bool:
+        Flag indicating if all the packages are installed."
+    """
+
+    for name in names:
+        try:
+            __import__(name)
+        except ImportError:
+            return False
+    return True