anemoi-datasets 0.5.26__py3-none-any.whl → 0.5.28__py3-none-any.whl

anemoi/datasets/create/sources/grib_index.py

@@ -10,10 +10,8 @@
 import logging
 import os
 import sqlite3
+from collections.abc import Iterator
 from typing import Any
-from typing import Iterator
-from typing import List
-from typing import Optional
 
 import earthkit.data as ekd
 import tqdm
@@ -21,7 +19,8 @@ from anemoi.transform.flavour import RuleBasedFlavour
 from cachetools import LRUCache
 from earthkit.data.indexing.fieldlist import FieldArray
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 LOG = logging.getLogger(__name__)
 
@@ -36,8 +35,8 @@ class GribIndex:
         self,
         database: str,
         *,
-        keys: Optional[List[str] | str] = None,
-        flavour: Optional[str] = None,
+        keys: list[str] | str | None = None,
+        flavour: str | None = None,
         update: bool = False,
         overwrite: bool = False,
     ) -> None:
@@ -157,7 +156,7 @@ class GribIndex:
         """Commit the current transaction to the database."""
         self.conn.commit()
 
-    def _get_metadata_keys(self) -> List[str]:
+    def _get_metadata_keys(self) -> list[str]:
         """Retrieve the metadata keys from the database.
 
         Returns
@@ -225,7 +224,7 @@ class GribIndex:
                 LOG.info(f"Path: {self._get_path(existing_record[1])}")
             raise
 
-    def _all_columns(self) -> List[str]:
+    def _all_columns(self) -> list[str]:
         """Retrieve all column names from the grib_index table.
 
         Returns
@@ -241,7 +240,7 @@ class GribIndex:
         self._columns = [col for col in columns if not col.startswith("_")]
         return self._columns
 
-    def _ensure_columns(self, columns: List[str]) -> None:
+    def _ensure_columns(self, columns: list[str]) -> None:
         """Add missing columns to the grib_index table.
 
         Parameters
@@ -324,7 +323,7 @@ class GribIndex:
 
         self._commit()
 
-    def _paramdb(self, category: int, discipline: int) -> Optional[dict]:
+    def _paramdb(self, category: int, discipline: int) -> dict | None:
         """Fetch parameter information from the parameter database.
 
         Parameters
@@ -355,7 +354,7 @@ class GribIndex:
         except Exception as e:
             LOG.warning(f"Failed to fetch information from parameter database: {e}")
 
-    def _param_grib2_info(self, paramId: int) -> List[dict]:
+    def _param_grib2_info(self, paramId: int) -> list[dict]:
         """Fetch GRIB2 parameter information for a given parameter ID.
 
         Parameters
@@ -383,7 +382,7 @@ class GribIndex:
             LOG.warning(f"Failed to fetch information from parameter database: {e}")
         return []
 
-    def _param_id_info(self, paramId: int) -> Optional[dict]:
+    def _param_id_info(self, paramId: int) -> dict | None:
         """Fetch detailed information for a given parameter ID.
 
         Parameters
@@ -412,7 +411,7 @@ class GribIndex:
 
         return None
 
-    def _param_id_unit(self, unitId: int) -> Optional[dict]:
+    def _param_id_unit(self, unitId: int) -> dict | None:
         """Fetch unit information for a given unit ID.
 
         Parameters
@@ -520,7 +519,7 @@ class GribIndex:
             raise ValueError(f"No path found for path_id {path_id}")
         return row[0]
 
-    def retrieve(self, dates: List[Any], **kwargs: Any) -> Iterator[Any]:
+    def retrieve(self, dates: list[Any], **kwargs: Any) -> Iterator[Any]:
         """Retrieve GRIB data from the database.
 
         Parameters
@@ -571,44 +570,47 @@ class GribIndex:
             yield data
 
 
-@legacy_source(__file__)
-def execute(
-    context: Any,
-    dates: List[Any],
-    indexdb: str,
-    flavour: Optional[str] = None,
-    **kwargs: Any,
-) -> FieldArray:
-    """Execute the GRIB data retrieval process.
-
-    Parameters
-    ----------
-    context : Any
-        The execution context.
-    dates : List[Any]
-        List of dates to retrieve data for.
-    indexdb : str
-        Path to the GRIB index database.
-    flavour : Optional[str], optional
-        Flavour configuration for mapping fields, by default None.
-    **kwargs : Any
-        Additional filtering criteria.
-
-    Returns
-    -------
-    FieldArray
-        An array of retrieved GRIB fields.
-    """
-    index = GribIndex(indexdb)
-    result = []
-
-    if flavour is not None:
-        flavour = RuleBasedFlavour(flavour)
-
-    for grib in index.retrieve(dates, **kwargs):
-        field = ekd.from_source("memory", grib)[0]
-        if flavour:
-            field = flavour.apply(field)
-        result.append(field)
-
-    return FieldArray(result)
+@source_registry.register("grib_index")
+class GribIndexSource(LegacySource):
+
+    @staticmethod
+    def _execute(
+        context: Any,
+        dates: list[Any],
+        indexdb: str,
+        flavour: str | None = None,
+        **kwargs: Any,
+    ) -> FieldArray:
+        """Execute the GRIB data retrieval process.
+
+        Parameters
+        ----------
+        context : Any
+            The execution context.
+        dates : List[Any]
+            List of dates to retrieve data for.
+        indexdb : str
+            Path to the GRIB index database.
+        flavour : Optional[str], optional
+            Flavour configuration for mapping fields, by default None.
+        **kwargs : Any
+            Additional filtering criteria.
+
+        Returns
+        -------
+        FieldArray
+            An array of retrieved GRIB fields.
+        """
+        index = GribIndex(indexdb)
+        result = []
+
+        if flavour is not None:
+            flavour = RuleBasedFlavour(flavour)
+
+        for grib in index.retrieve(dates, **kwargs):
+            field = ekd.from_source("memory", grib)[0]
+            if flavour:
+                field = flavour.apply(field)
+            result.append(field)
+
+        return FieldArray(result)

anemoi/datasets/create/sources/hindcasts.py

@@ -9,20 +9,18 @@
 
 import logging
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Union
 
 from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.create.sources.mars import mars
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 LOGGER = logging.getLogger(__name__)
 
 
-def _to_list(x: Union[list, tuple, Any]) -> List[Any]:
+def _to_list(x: list | tuple | Any) -> list[Any]:
     """Converts the input to a list if it is not already a list or tuple.
 
     Parameters
@@ -40,57 +38,57 @@ def _to_list(x: Union[list, tuple, Any]) -> List[Any]:
     return [x]
 
 
-@legacy_source(__file__)
-def hindcasts(context: Any, dates: List[Any], **request: Dict[str, Any]) -> MultiFieldList:
-    """Generates hindcast requests based on the provided dates and request parameters.
-
-    Parameters
-    ----------
-    context : Any
-        The context containing the dates provider and trace method.
-    dates : List[Any]
-        A list of dates for which to generate hindcast requests.
-    request : Dict[str, Any]
-        Additional request parameters.
-
-    Returns
-    -------
-    MultiFieldList
-        A MultiFieldList containing the hindcast data.
-    """
-    from anemoi.datasets.dates import HindcastsDates
-
-    provider = context.dates_provider
-    assert isinstance(provider, HindcastsDates)
-
-    context.trace("H️", f"hindcasts {len(dates)=}")
-
-    request["param"] = _to_list(request["param"])
-    request["step"] = _to_list(request.get("step", 0))
-    request["step"] = [int(_) for _ in request["step"]]
-
-    context.trace("H️", f"hindcast {request}")
-
-    requests = []
-    for d in dates:
-        r = request.copy()
-        hindcast = provider.mapping[d]
-        r["hdate"] = hindcast.hdate.strftime("%Y-%m-%d")
-        r["date"] = hindcast.refdate.strftime("%Y-%m-%d")
-        r["time"] = hindcast.refdate.strftime("%H")
-        r["step"] = hindcast.step
-        requests.append(r)
-
-    if len(requests) == 0:
-        return MultiFieldList([])
-
-    return mars(
-        context,
-        dates,
-        *requests,
-        date_key="hdate",
-        request_already_using_valid_datetime=True,
-    )
-
-
-execute = hindcasts
+@source_registry.register("hindcasts")
+class HindcastsSource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[Any], **request: dict[str, Any]) -> MultiFieldList:
+        """Generates hindcast requests based on the provided dates and request parameters.
+
+        Parameters
+        ----------
+        context : Any
+            The context containing the dates provider and trace method.
+        dates : List[Any]
+            A list of dates for which to generate hindcast requests.
+        request : Dict[str, Any]
+            Additional request parameters.
+
+        Returns
+        -------
+        MultiFieldList
+            A MultiFieldList containing the hindcast data.
+        """
+        from anemoi.datasets.dates import HindcastsDates
+
+        provider = context.dates_provider
+        assert isinstance(provider, HindcastsDates)
+
+        context.trace("H️", f"hindcasts {len(dates)=}")
+
+        request["param"] = _to_list(request["param"])
+        request["step"] = _to_list(request.get("step", 0))
+        request["step"] = [int(_) for _ in request["step"]]
+
+        context.trace("H️", f"hindcast {request}")
+
+        requests = []
+        for d in dates:
+            r = request.copy()
+            hindcast = provider.mapping[d]
+            r["hdate"] = hindcast.hdate.strftime("%Y-%m-%d")
+            r["date"] = hindcast.refdate.strftime("%Y-%m-%d")
+            r["time"] = hindcast.refdate.strftime("%H")
+            r["step"] = hindcast.step
+            requests.append(r)
+
+        if len(requests) == 0:
+            return MultiFieldList([])
+
+        return mars(
+            context,
+            dates,
+            *requests,
+            date_key="hdate",
+            request_already_using_valid_datetime=True,
+        )

anemoi/datasets/create/sources/legacy.py

@@ -8,16 +8,13 @@
 # nor does it submit to any jurisdiction.
 
 
-import inspect
 import logging
-import os
+from abc import abstractmethod
 from typing import Any
-from typing import Callable
 
-from anemoi.datasets.create.input.template import resolve
+from anemoi.datasets.create.input.context import Context
 
 from ..source import Source
-from . import source_registry
 
 LOG = logging.getLogger(__name__)
 
@@ -27,7 +24,7 @@ class LegacySource(Source):
 
     Parameters
     ----------
-    context : Any
+    context : Context
         The context in which the source is created.
     *args : tuple
         Positional arguments.
@@ -35,64 +32,15 @@ class LegacySource(Source):
         Keyword arguments.
     """
 
-    def __init__(self, context: Any, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, context: Context, *args: Any, **kwargs: Any) -> None:
         super().__init__(context, *args, **kwargs)
         self.args = args
         self.kwargs = kwargs
 
+    @staticmethod
+    @abstractmethod
+    def _execute(context, *args, **kwargs):
+        pass
 
-class legacy_source:
-    """A decorator class for legacy sources.
-
-    Parameters
-    ----------
-    name : str
-        The name of the legacy source.
-    """
-
-    def __init__(self, name: str) -> None:
-        name, _ = os.path.splitext(os.path.basename(name))
-        self.name = name
-
-    def __call__(self, execute: Callable) -> Callable:
-        """Call method to wrap the execute function.
-
-        Parameters
-        ----------
-        execute : function
-            The execute function to be wrapped.
-
-        Returns
-        -------
-        function
-            The wrapped execute function.
-        """
-        this = self
-        name = f"Legacy{self.name.title()}Source"
-        source = ".".join([execute.__module__, execute.__name__])
-
-        def execute_wrapper(self, dates) -> Any:
-            """Wrapper method to call the execute function."""
-
-            args, kwargs = resolve(self.context, (self.args, self.kwargs))
-
-            try:
-                return execute(self.context, dates, *args, **kwargs)
-            except TypeError:
-                LOG.error(f"Error executing source {this.name} from {source}")
-                LOG.error(f"Function signature is: {inspect.signature(execute)}")
-                LOG.error(f"Arguments are: {args=}, {kwargs=}")
-                raise
-
-        klass = type(
-            name,
-            (LegacySource,),
-            {
-                "execute": execute_wrapper,
-                "_source": source,
-            },
-        )
-
-        source_registry.register(self.name)(klass)
-
-        return execute
+    def execute(self, dates: Any) -> Any:
+        return self._execute(self.context, dates, *self.args, **self.kwargs)