anemoi-datasets 0.5.15__py3-none-any.whl → 0.5.17__py3-none-any.whl

anemoi/datasets/create/input/function.py

@@ -9,8 +9,12 @@
 
 import logging
 from functools import cached_property
+from typing import Any
+from typing import Dict
 
-from ..functions import import_function
+from earthkit.data import FieldList
+
+from ...dates.groups import GroupOfDates
 from .action import Action
 from .misc import _tidy
 from .misc import assert_fieldlist
@@ -30,87 +34,200 @@ class FunctionContext:
     to the functions from the other actions and filters and results.
     """
 
-    def __init__(self, owner):
-        self.owner = owner
-        self.use_grib_paramid = owner.context.use_grib_paramid
+    def __init__(self, owner: Result) -> None:
+        """Initializes a FunctionContext instance.
 
-    def trace(self, emoji, *args):
+        Parameters
+        ----------
+        owner : object
+            The owner object.
+        """
+        self.owner = owner
+        self.use_grib_paramid: bool = owner.context.use_grib_paramid
+
+    def trace(self, emoji: str, *args: Any) -> None:
+        """Traces the given arguments with an emoji.
+
+        Parameters
+        ----------
+        emoji : str
+            The emoji to use.
+        *args : Any
+            The arguments to trace.
+        """
         trace(emoji, *args)
 
-    def info(self, *args, **kwargs):
+    def info(self, *args: Any, **kwargs: Any) -> None:
+        """Logs an info message.
+
+        Parameters
+        ----------
+        *args : Any
+            The arguments for the log message.
+        **kwargs : Any
+            The keyword arguments for the log message.
+        """
         LOG.info(*args, **kwargs)
 
     @property
-    def dates_provider(self):
+    def dates_provider(self) -> object:
+        """Returns the dates provider."""
         return self.owner.group_of_dates.provider
 
     @property
-    def partial_ok(self):
+    def partial_ok(self) -> bool:
+        """Returns whether partial results are acceptable."""
         return self.owner.group_of_dates.partial_ok
 
 
 class FunctionAction(Action):
-    def __init__(self, context, action_path, _name, **kwargs):
+    """Represents an action that executes a function.
+
+    Attributes
+    ----------
+    name : str
+        The name of the function.
+    """
+
+    def __init__(self, context: object, action_path: list, _name: str, source, **kwargs: Dict[str, Any]) -> None:
+        """Initializes a FunctionAction instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : list
+            The action path.
+        _name : str
+            The name of the function.
+        **kwargs : Dict[str, Any]
+            Additional keyword arguments.
+        """
         super().__init__(context, action_path, **kwargs)
-        self.name = _name
+        self.name: str = _name
+        self.source = source
 
     @trace_select
-    def select(self, group_of_dates):
+    def select(self, group_of_dates: GroupOfDates) -> "FunctionResult":
+        """Selects the function result for the given group of dates.
+
+        Parameters
+        ----------
+        group_of_dates : GroupOfDates
+            The group of dates.
+
+        Returns
+        -------
+        FunctionResult
+            The function result instance.
+        """
         return FunctionResult(self.context, self.action_path, group_of_dates, action=self)
 
-    @property
-    def function(self):
-        # name, delta = parse_function_name(self.name)
-        return import_function(self.name, "sources")
-
-    def __repr__(self):
-        content = ""
+    def __repr__(self) -> str:
+        """Returns a string representation of the FunctionAction instance."""
+        content: str = ""
         content += ",".join([self._short_str(a) for a in self.args])
         content += " ".join([self._short_str(f"{k}={v}") for k, v in self.kwargs.items()])
         content = self._short_str(content)
-        return super().__repr__(_inline_=content, _indent_=" ")
+        return self._repr(_inline_=content, _indent_=" ")
+
+    def _trace_select(self, group_of_dates: GroupOfDates) -> str:
+        """Traces the selection of the function for the given group of dates.
 
-    def _trace_select(self, group_of_dates):
+        Parameters
+        ----------
+        group_of_dates : GroupOfDates
+            The group of dates.
+
+        Returns
+        -------
+        str
+            The trace string.
+        """
         return f"{self.name}({group_of_dates})"
 
 
 class FunctionResult(Result):
-    def __init__(self, context, action_path, group_of_dates, action):
+    """Represents the result of executing a function.
+
+    Attributes
+    ----------
+    action : Action
+        The action instance.
+    args : tuple
+        The positional arguments for the function.
+    kwargs : dict
+        The keyword arguments for the function.
+    """
+
+    def __init__(self, context: object, action_path: list, group_of_dates: GroupOfDates, action: Action) -> None:
+        """Initializes a FunctionResult instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : list
+            The action path.
+        group_of_dates : GroupOfDates
+            The group of dates.
+        action : Action
+            The action instance.
+        """
         super().__init__(context, action_path, group_of_dates)
         assert isinstance(action, Action), type(action)
-        self.action = action
+        self.action: Action = action
 
         self.args, self.kwargs = substitute(context, (self.action.args, self.action.kwargs))
 
-    def _trace_datasource(self, *args, **kwargs):
+    def _trace_datasource(self, *args: Any, **kwargs: Any) -> str:
+        """Traces the datasource for the given arguments.
+
+        Parameters
+        ----------
+        *args : Any
+            The arguments.
+        **kwargs : Any
+            The keyword arguments.
+
+        Returns
+        -------
+        str
+            The trace string.
+        """
         return f"{self.action.name}({self.group_of_dates})"
 
     @cached_property
     @assert_fieldlist
     @notify_result
     @trace_datasource
-    def datasource(self):
+    def datasource(self) -> FieldList:
+        """Returns the datasource for the function result."""
         args, kwargs = resolve(self.context, (self.args, self.kwargs))
+        self.action.source.context = FunctionContext(self)
 
-        try:
-            return _tidy(
-                self.action.function(
-                    FunctionContext(self),
-                    list(self.group_of_dates),  # Will provide a list of datetime objects
-                    *args,
-                    **kwargs,
-                )
+        return _tidy(
+            self.action.source.execute(
+                self.group_of_dates,  # Will provide a list of datetime objects
+                *args,
+                **kwargs,
             )
-        except Exception:
-            LOG.error(f"Error in {self.action.function.__name__}", exc_info=True)
-            raise
+        )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Returns a string representation of the FunctionResult instance."""
         try:
             return f"{self.action.name}({self.group_of_dates})"
         except Exception:
             return f"{self.__class__.__name__}(unitialised)"
 
     @property
-    def function(self):
+    def function(self) -> None:
+        """Raises NotImplementedError as this property is not implemented.
+
+        Raises
+        ------
+        NotImplementedError
+            Always raised.
+        """
         raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")

anemoi/datasets/create/input/join.py

@@ -9,7 +9,12 @@
 
 import logging
 from functools import cached_property
+from typing import Any
+from typing import List
 
+from earthkit.data import FieldList
+
+from ...dates.groups import GroupOfDates
 from .action import Action
 from .action import action_factory
 from .empty import EmptyResult
@@ -24,35 +29,102 @@ LOG = logging.getLogger(__name__)
 
 
 class JoinResult(Result):
-    def __init__(self, context, action_path, group_of_dates, results, **kwargs):
+    """Represents a result that combines multiple results.
+
+    Attributes
+    ----------
+    context : object
+        The context object.
+    action_path : list
+        The action path.
+    group_of_dates : GroupOfDates
+        The group of dates.
+    results : List[Result]
+        The list of results.
+    """
+
+    def __init__(
+        self, context: object, action_path: list, group_of_dates: GroupOfDates, results: List[Result], **kwargs: Any
+    ) -> None:
+        """Initializes a JoinResult instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : list
+            The action path.
+        group_of_dates : GroupOfDates
+            The group of dates.
+        results : List[Result]
+            The list of results.
+        """
         super().__init__(context, action_path, group_of_dates)
-        self.results = [r for r in results if not r.empty]
+        self.results: List[Result] = [r for r in results if not r.empty]
 
     @cached_property
     @assert_fieldlist
     @notify_result
     @trace_datasource
-    def datasource(self):
-        ds = EmptyResult(self.context, self.action_path, self.group_of_dates).datasource
+    def datasource(self) -> FieldList:
+        """Returns the combined datasource from all results."""
+        ds: FieldList = EmptyResult(self.context, self.action_path, self.group_of_dates).datasource
         for i in self.results:
             ds += i.datasource
         return _tidy(ds)
 
-    def __repr__(self):
-        content = "\n".join([str(i) for i in self.results])
-        return super().__repr__(content)
+    def __repr__(self) -> str:
+        """Returns a string representation of the JoinResult instance."""
+        content: str = "\n".join([str(i) for i in self.results])
+        return self._repr(content)
 
 
 class JoinAction(Action):
-    def __init__(self, context, action_path, *configs):
+    """Represents an action that combines multiple actions.
+
+    Attributes
+    ----------
+    context : object
+        The context object.
+    action_path : list
+        The action path.
+    actions : List[Action]
+        The list of actions.
+    """
+
+    def __init__(self, context: object, action_path: list, *configs: dict) -> None:
+        """Initializes a JoinAction instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : list
+            The action path.
+        *configs : dict
+            The configuration dictionaries.
+        """
         super().__init__(context, action_path, *configs)
-        self.actions = [action_factory(c, context, action_path + [str(i)]) for i, c in enumerate(configs)]
+        self.actions: List[Action] = [action_factory(c, context, action_path + [str(i)]) for i, c in enumerate(configs)]
 
-    def __repr__(self):
-        content = "\n".join([str(i) for i in self.actions])
-        return super().__repr__(content)
+    def __repr__(self) -> str:
+        """Returns a string representation of the JoinAction instance."""
+        content: str = "\n".join([str(i) for i in self.actions])
+        return self._repr(content)
 
     @trace_select
-    def select(self, group_of_dates):
-        results = [a.select(group_of_dates) for a in self.actions]
+    def select(self, group_of_dates: GroupOfDates) -> JoinResult:
+        """Selects the results for the given group of dates.
+
+        Parameters
+        ----------
+        group_of_dates : GroupOfDates
+            The group of dates.
+
+        Returns
+        -------
+        JoinResult
+            The combined result for the given group of dates.
+        """
+        results: List[Result] = [a.select(group_of_dates) for a in self.actions]
         return JoinResult(self.context, self.action_path, group_of_dates, results)

anemoi/datasets/create/input/misc.py

@@ -9,17 +9,30 @@
 
 import logging
 from functools import wraps
+from typing import Any
+from typing import Callable
+from typing import Tuple
+from typing import Union
 
+from earthkit.data import FieldList
 from earthkit.data.core.fieldlist import MultiFieldList
-from earthkit.data.indexing.fieldlist import FieldList
-
-from ..functions import import_function
 
 LOG = logging.getLogger(__name__)
 
 
-def parse_function_name(name):
+def parse_function_name(name: str) -> Tuple[str, Union[int, None]]:
+    """Parses a function name to extract the base name and an optional time delta.
+
+    Parameters
+    ----------
+    name : str
+        The function name to parse.
 
+    Returns
+    -------
+    tuple of (str, int or None)
+        The base name and an optional time delta.
+    """
     if name.endswith("h") and name[:-1].isdigit():
 
         if "-" in name:
@@ -40,19 +53,23 @@ def parse_function_name(name):
     return name, None
 
 
-def is_function(name, kind):
-    name, _ = parse_function_name(name)
-    try:
-        import_function(name, kind)
-        return True
-    except ImportError as e:
-        print(e)
-        return False
+def assert_fieldlist(method: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator to assert that the result of a method is an instance of FieldList.
+
+    Parameters
+    ----------
+    method : Callable[..., Any]
+        The method to decorate.
 
+    Returns
+    -------
+    Callable[..., Any]
+        The decorated method.
+    """
 
-def assert_fieldlist(method):
     @wraps(method)
-    def wrapper(self, *args, **kwargs):
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+
         result = method(self, *args, **kwargs)
         assert isinstance(result, FieldList), type(result)
         return result
@@ -60,17 +77,50 @@ def assert_fieldlist(method):
     return wrapper
 
 
-def assert_is_fieldlist(obj):
+def assert_is_fieldlist(obj: object) -> None:
+    """Asserts that the given object is an instance of FieldList.
+
+    Parameters
+    ----------
+    obj : object
+        The object to check.
+    """
     assert isinstance(obj, FieldList), type(obj)
 
 
-def _flatten(ds):
+def _flatten(ds: Union[MultiFieldList, FieldList]) -> list:
+    """Flattens a MultiFieldList or FieldList into a list of FieldList objects.
+
+    Parameters
+    ----------
+    ds : Union[MultiFieldList, FieldList]
+        The dataset to flatten.
+
+    Returns
+    -------
+    list
+        A list of FieldList objects.
+    """
     if isinstance(ds, MultiFieldList):
         return [_tidy(f) for s in ds._indexes for f in _flatten(s)]
     return [ds]
 
 
-def _tidy(ds, indent=0):
+def _tidy(ds: Union[MultiFieldList, FieldList], indent: int = 0) -> Union[MultiFieldList, FieldList]:
+    """Tidies up a MultiFieldList or FieldList by removing empty sources.
+
+    Parameters
+    ----------
+    ds : Union[MultiFieldList, FieldList]
+        The dataset to tidy.
+    indent : int, optional
+        The indentation level. Defaults to 0.
+
+    Returns
+    -------
+    Union[MultiFieldList, FieldList]
+        The tidied dataset.
+    """
     if isinstance(ds, MultiFieldList):
 
         sources = [s for s in _flatten(ds) if len(s) > 0]

anemoi/datasets/create/input/pipe.py

@@ -9,6 +9,7 @@
 
 import json
 import logging
+from typing import Any
 
 from .action import Action
 from .action import action_factory
@@ -19,21 +20,47 @@ LOG = logging.getLogger(__name__)
 
 
 class PipeAction(Action):
-    def __init__(self, context, action_path, *configs):
+    """A class to represent a pipeline of actions."""
+
+    def __init__(self, context: Any, action_path: list, *configs: dict) -> None:
+        """Initialize the PipeAction.
+
+        Parameters
+        ----------
+        context : Any
+            The context for the action.
+        action_path : list
+            The path of the action.
+        configs : dict
+            The configurations for the actions.
+        """
         super().__init__(context, action_path, *configs)
         if len(configs) <= 1:
             raise ValueError(
                 f"PipeAction requires at least two actions, got {len(configs)}\n{json.dumps(configs, indent=2)}"
             )
 
-        current = action_factory(configs[0], context, action_path + ["0"])
+        current: Any = action_factory(configs[0], context, action_path + ["0"])
         for i, c in enumerate(configs[1:]):
             current = step_factory(c, context, action_path + [str(i + 1)], previous_step=current)
-        self.last_step = current
+        self.last_step: Any = current
 
     @trace_select
-    def select(self, group_of_dates):
+    def select(self, group_of_dates: Any) -> Any:
+        """Select data based on the group of dates.
+
+        Parameters
+        ----------
+        group_of_dates : Any
+            The group of dates to select data for.
+
+        Returns
+        -------
+        Any
+            The selected data.
+        """
         return self.last_step.select(group_of_dates)
 
-    def __repr__(self):
-        return super().__repr__(self.last_step)
+    def __repr__(self) -> str:
+        """Return a string representation of the PipeAction."""
+        return f"PipeAction({self.last_step})"