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

anemoi/utils/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.13'
-__version_tuple__ = version_tuple = (0, 4, 13)
+__version__ = version = '0.4.14'
+__version_tuple__ = version_tuple = (0, 4, 14)

anemoi/utils/provenance.py

@@ -10,10 +10,10 @@
 
 """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
- - ...
+- 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

anemoi/utils/registry.py

@@ -12,9 +12,12 @@ import importlib
 import logging
 import os
 import sys
+import warnings
+from functools import cached_property
 from typing import Any
 from typing import Callable
 from typing import Dict
+from typing import List
 from typing import Optional
 from typing import Union
 
@@ -22,6 +25,8 @@ import entrypoints
 
 LOG = logging.getLogger(__name__)
 
+DEBUG_ANEMOI_REGISTRY = int(os.environ.get("DEBUG_ANEMOI_REGISTRY", "0"))
+
 
 class Wrapper:
     """A wrapper for the registry.
@@ -55,6 +60,22 @@ class Wrapper:
         return factory
 
 
+class Error:
+    """An error class. Used in place of a plugin that failed to load.
+
+    Parameters
+    ----------
+    error : Exception
+        The error.
+    """
+
+    def __init__(self, error: Exception):
+        self.error = error
+
+    def __call__(self, *args, **kwargs):
+        raise self.error
+
+
 _BY_KIND = {}
 
 
@@ -71,7 +92,8 @@ class Registry:
 
     def __init__(self, package: str, key: str = "_type"):
         self.package = package
-        self.registered = {}
+        self.__registered = {}
+        self._sources = {}
         self.kind = package.split(".")[-1]
         self.key = key
         _BY_KIND[self.kind] = self
@@ -92,7 +114,9 @@ class Registry:
         """
         return _BY_KIND.get(kind)
 
-    def register(self, name: str, factory: Optional[Callable] = None) -> Optional[Wrapper]:
+    def register(
+        self, name: str, factory: Optional[Callable] = None, source: Optional[Any] = None
+    ) -> Optional[Wrapper]:
         """Register a factory with the registry.
 
         Parameters
@@ -101,6 +125,8 @@ class Registry:
             The name of the factory.
         factory : Callable, optional
             The factory to register, by default None.
+        source : Any, optional
+            The source of the factory, by default None.
 
         Returns
         -------
@@ -108,12 +134,19 @@ class Registry:
             A wrapper if the factory is None, otherwise None.
         """
         if factory is None:
+            # This happens when the @register decorator is used
             return Wrapper(name, self)
 
-        self.registered[name] = factory
+        if source is None:
+            source = getattr(factory, "_source") if hasattr(factory, "_source") else factory
+
+        if name in self.__registered:
+            warnings.warn(f"Factory '{name}' is already registered in {self.package}")
+            warnings.warn(f"Existing: {self._sources[name]}")
+            warnings.warn(f"New: {source}")
 
-    # def registered(self, name: str):
-    #     return name in self.registered
+        self.__registered[name] = factory
+        self._sources[name] = source
 
     def _load(self, file: str) -> None:
         """Load a module from a file.
@@ -126,8 +159,30 @@ class Registry:
         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)
+        except Exception as e:
+            if DEBUG_ANEMOI_REGISTRY:
+                raise
+            self._registered[name] = Error(e)
+
+    def is_registered(self, name: str) -> bool:
+        """Check if a factory is registered.
+
+        Parameters
+        ----------
+        name : str
+            The name of the factory.
+
+        Returns
+        -------
+        bool
+            Whether the factory is registered.
+        """
+        ok = name in self.factories
+        if not ok:
+            LOG.error(f"Cannot find '{name}' in {self.package}")
+            for e in self.factories:
+                LOG.info(f"Registered: {e} ({self._sources.get(e)})")
+        return ok
 
     def lookup(self, name: str, *, return_none: bool = False) -> Optional[Callable]:
         """Lookup a factory by name.
@@ -144,9 +199,22 @@ class Registry:
         Callable, optional
             The factory if found, otherwise None.
         """
-        # print('✅✅✅✅✅✅✅✅✅✅✅✅✅', name, self.registered)
-        if name in self.registered:
-            return self.registered[name]
+        if return_none:
+            return self.factories.get(name)
+
+        factory = self.factories.get(name)
+        if factory is None:
+
+            LOG.error(f"Cannot find '{name}' in {self.package}")
+            for e in self.factories:
+                LOG.info(f"Registered: {e} ({self._sources.get(e)})")
+
+            raise ValueError(f"Cannot find '{name}' in {self.package}")
+
+        return factory
+
+    @cached_property
+    def factories(self) -> Dict[str, Callable]:
 
         directory = sys.modules[self.package].__path__[0]
 
@@ -167,25 +235,41 @@ class Registry:
             if file.endswith(".py"):
                 self._load(file)
 
-        entrypoint_group = f"anemoi.{self.kind}"
-        for entry_point in entrypoints.get_group_all(entrypoint_group):
-            if entry_point.name == name:
-                if name in self.registered:
-                    LOG.warning(
-                        f"Overwriting builtin '{name}' from {self.package} with plugin '{entry_point.module_name}'"
-                    )
-                self.registered[name] = entry_point.load()
+        bits = self.package.split(".")
+        # We assume a name like anemoi.datasets.create.sources, with kind = sources
+        assert bits[-1] == self.kind, (self.package, self.kind)
+        assert len(bits) > 1, self.package
+
+        groups = []
+        middle = bits[1:-1]
+        while True:
+            group = ".".join([bits[0], *middle, bits[-1]])
+            groups.append(group)
+            if len(middle) == 0:
+                break
+            middle.pop()
 
-        if name not in self.registered:
-            if return_none:
-                return None
+        groups.reverse()
 
-            for e in self.registered:
-                LOG.info(f"Registered: {e}")
+        LOG.debug("Loading plugins from %s", groups)
 
-            raise ValueError(f"Cannot load '{name}' from {self.package}")
+        for entrypoint_group in groups:
+            for entry_point in entrypoints.get_group_all(entrypoint_group):
+                source = entry_point.distro
+                try:
+                    self.register(entry_point.name, entry_point.load(), source=source)
+                except Exception as e:
+                    if DEBUG_ANEMOI_REGISTRY:
+                        raise
+                    self.register(entry_point.name, Error(e), source=source)
 
-        return self.registered[name]
+        return self.__registered
+
+    @property
+    def registered(self) -> List[str]:
+        """Get the registered factories."""
+
+        return sorted(self.factories.keys())
 
     def create(self, name: str, *args: Any, **kwargs: Any) -> Any:
         """Create an instance using a factory.
@@ -207,9 +291,6 @@ class Registry:
         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: Union[str, Dict[str, Any]], *args: Any, **kwargs: Any) -> Any:
         """Create an instance from a configuration.
 

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

{anemoi_utils-0.4.13.dist-info → anemoi_utils-0.4.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.13
+Version: 0.4.14
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License
@@ -256,3 +256,4 @@ Requires-Dist: pytest; extra == "tests"
 Provides-Extra: text
 Requires-Dist: termcolor; extra == "text"
 Requires-Dist: wcwidth; extra == "text"
+Dynamic: license-file

{anemoi_utils-0.4.13.dist-info → anemoi_utils-0.4.14.dist-info}/RECORD

@@ -1,6 +1,6 @@
 anemoi/utils/__init__.py,sha256=uVhpF-VjIl_4mMywOVtgTutgsdIsqz-xdkwxeMhzuag,730
 anemoi/utils/__main__.py,sha256=6LlE4MYrPvqqrykxXh7XMi50UZteUY59NeM8P9Zs2dU,910
-anemoi/utils/_version.py,sha256=NEWZMsf9hAPHbHDxsYo2Rgs87lmnioXpPKJ62lRxhqs,513
+anemoi/utils/_version.py,sha256=ggKLQlA5gegh3l4LAMWuwHPmL8zgZ4b0pYtFx7quk78,513
 anemoi/utils/caching.py,sha256=rXbeAmpBcMbbfN4EVblaHWKicsrtx1otER84FEBtz98,6183
 anemoi/utils/checkpoints.py,sha256=N4WpAZXa4etrpSEKhHqUUtG2-x9w3FJMHcLO-dDAXPY,9600
 anemoi/utils/cli.py,sha256=IyZfnSw0u0yYnrjOrzvm2RuuKvDk4cVb8pf8BkaChgA,6209
@@ -13,11 +13,12 @@ anemoi/utils/grids.py,sha256=edTrMK8hpE9ZBzSfwcRftgk0jljNAK3i8CraadILQoM,4427
 anemoi/utils/hindcasts.py,sha256=iYVIxSNFL2HJcc_k1abCFLkpJFGHT8WKRIR4wcAwA3s,2144
 anemoi/utils/humanize.py,sha256=hCrHr5ppREuJR-tBqRqynqe58BHR6Ga_gCQqgEmmrfU,25301
 anemoi/utils/logs.py,sha256=naTgrmPwWHD4eekFttXftS4gtcAGYHpCqG4iwYprNDA,1804
-anemoi/utils/provenance.py,sha256=l6TRVadM5l3SKUvYM20EQn0TbolSY6vbbZ_WqekjxwM,14619
-anemoi/utils/registry.py,sha256=J2onjVqR4LLYRAj5sioobsdxo_e4kxuBOSiNeAQbGe8,7061
+anemoi/utils/provenance.py,sha256=tIIgweS0EJyaYzgKwuv3iWny-Gz7N5e5CNWH5MeSYWU,14615
+anemoi/utils/registry.py,sha256=zTnVSCecmrI6SkZhF4ipV7WKiZYHEBJ_5ZwqhRGM4T0,9287
 anemoi/utils/s3.py,sha256=xMT48kbcelcjjqsaU567WI3oZ5eqo88Rlgyx5ECszAU,4074
 anemoi/utils/sanitise.py,sha256=ZYGdSX6qihQANr3pHZjbKnoapnzP1KcrWdW1Ul1mOGk,3668
 anemoi/utils/sanitize.py,sha256=43ZKDcfVpeXSsJ9TFEc9aZnD6oe2cUh151XnDspM98M,462
+anemoi/utils/testing.py,sha256=N1y4dfZLE9zqOhIR3o-933fdAdd9BxDvjcJx7SwFC9A,4803
 anemoi/utils/text.py,sha256=HkzIvi24obDceFLpJEwBJ9PmPrJUkQN2TrElJ-A87gU,14441
 anemoi/utils/timer.py,sha256=_leKMYza2faM7JKlGE7LCNy13rbdPnwaCF7PSrI_NmI,3895
 anemoi/utils/commands/__init__.py,sha256=5u_6EwdqYczIAgJfCwRSyQAYFEqh2ZuHHT57g9g7sdI,808
@@ -29,9 +30,9 @@ anemoi/utils/mars/requests.py,sha256=VFMHBVAAl0_2lOcMBa1lvaKHctN0lDJsI6_U4BucGew
 anemoi/utils/remote/__init__.py,sha256=-uaYFi4yRYFRf46ubQbJo86GCn6HE5VQrcaoyrmyW28,20704
 anemoi/utils/remote/s3.py,sha256=spQ8l0rwQjLZh9dZu5cOsYIvNwKihQfCJ6YsFYegeqI,17339
 anemoi/utils/remote/ssh.py,sha256=xNtsawh8okytCKRehkRCVExbHZj-CRUQNormEHglfuw,8088
-anemoi_utils-0.4.13.dist-info/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
-anemoi_utils-0.4.13.dist-info/METADATA,sha256=WSExhiI_8UePoAsCE4h65G2moeuMEgWZcgECBsisx30,15338
-anemoi_utils-0.4.13.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-anemoi_utils-0.4.13.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
-anemoi_utils-0.4.13.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_utils-0.4.13.dist-info/RECORD,,
+anemoi_utils-0.4.14.dist-info/licenses/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
+anemoi_utils-0.4.14.dist-info/METADATA,sha256=VLuh3lxCAuACOrJRnE_Dm0zX_pVfIHoppmRrP505pkY,15360
+anemoi_utils-0.4.14.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+anemoi_utils-0.4.14.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
+anemoi_utils-0.4.14.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_utils-0.4.14.dist-info/RECORD,,

{anemoi_utils-0.4.13.dist-info → anemoi_utils-0.4.14.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (76.0.0)
+Generator: setuptools (77.0.3)
 Root-Is-Purelib: true
 Tag: py3-none-any