anemoi-datasets 0.5.16__py3-none-any.whl → 0.5.18__py3-none-any.whl

anemoi/datasets/create/input/concat.py

@@ -10,9 +10,16 @@
 import logging
 from copy import deepcopy
 from functools import cached_property
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Union
+
+from earthkit.data import FieldList
 
 from anemoi.datasets.dates import DatesProvider
 
+from ...dates.groups import GroupOfDates
 from .action import Action
 from .action import action_factory
 from .empty import EmptyResult
@@ -27,7 +34,31 @@ LOG = logging.getLogger(__name__)
 
 
 class ConcatResult(Result):
-    def __init__(self, context, action_path, group_of_dates, results, **kwargs):
+    """Represents the result of concatenating multiple results."""
+
+    def __init__(
+        self,
+        context: object,
+        action_path: List[str],
+        group_of_dates: GroupOfDates,
+        results: List[Result],
+        **kwargs: Any,
+    ) -> None:
+        """Initializes a ConcatResult instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : List[str]
+            The action path.
+        group_of_dates : GroupOfDates
+            The group of dates.
+        results : List[Result]
+            The list of results.
+        kwargs : Any
+            Additional keyword arguments.
+        """
         super().__init__(context, action_path, group_of_dates)
         self.results = [r for r in results if not r.empty]
 
@@ -35,15 +66,16 @@ class ConcatResult(Result):
     @assert_fieldlist
     @notify_result
     @trace_datasource
-    def datasource(self):
+    def datasource(self) -> FieldList:
+        """Returns the concatenated datasource from all results."""
         ds = EmptyResult(self.context, self.action_path, self.group_of_dates).datasource
         for i in self.results:
             ds += i.datasource
         return _tidy(ds)
 
     @property
-    def variables(self):
-        """Check that all the results objects have the same variables."""
+    def variables(self) -> List[str]:
+        """Returns the list of variables, ensuring all results have the same variables."""
         variables = None
         for f in self.results:
             if f.empty:
@@ -54,13 +86,33 @@ class ConcatResult(Result):
         assert variables is not None, self.results
         return variables
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Returns a string representation of the ConcatResult instance.
+
+        Returns
+        -------
+        str
+            A string representation of the ConcatResult instance.
+        """
         content = "\n".join([str(i) for i in self.results])
-        return super().__repr__(content)
+        return self._repr(content)
 
 
 class ConcatAction(Action):
-    def __init__(self, context, action_path, *configs):
+    """Represents an action that concatenates multiple actions based on their dates."""
+
+    def __init__(self, context: object, action_path: List[str], *configs: Dict[str, Any]) -> None:
+        """Initializes a ConcatAction instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : List[str]
+            The action path.
+        configs : Dict[str, Any]
+            The configuration dictionaries.
+        """
         super().__init__(context, action_path, *configs)
         parts = []
         for i, cfg in enumerate(configs):
@@ -74,12 +126,31 @@ class ConcatAction(Action):
             parts.append((filtering_dates, action))
         self.parts = parts
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Returns a string representation of the ConcatAction instance.
+
+        Returns
+        -------
+        str
+            A string representation of the ConcatAction instance.
+        """
         content = "\n".join([str(i) for i in self.parts])
-        return super().__repr__(content)
+        return self._repr(content)
 
     @trace_select
-    def select(self, group_of_dates):
+    def select(self, group_of_dates: GroupOfDates) -> Union[ConcatResult, EmptyResult]:
+        """Selects the concatenated result for the given group of dates.
+
+        Parameters
+        ----------
+        group_of_dates : GroupOfDates
+            The group of dates.
+
+        Returns
+        -------
+        Union[ConcatResult, EmptyResult]
+            The concatenated result or an empty result.
+        """
         from anemoi.datasets.dates.groups import GroupOfDates
 
         results = []

anemoi/datasets/create/input/context.py

@@ -9,6 +9,10 @@
 
 import logging
 import textwrap
+from typing import Any
+from typing import List
+from typing import Tuple
+from typing import Union
 
 from anemoi.utils.humanize import plural
 
@@ -19,18 +23,37 @@ LOG = logging.getLogger(__name__)
 
 
 class Context:
-    def __init__(self):
+    """Class to handle the build context in the dataset creation process."""
+
+    def __init__(self) -> None:
+        """Initializes a Context instance."""
         # used_references is a set of reference paths that will be needed
         self.used_references = set()
         # results is a dictionary of reference path -> obj
         self.results = {}
 
-    def will_need_reference(self, key):
+    def will_need_reference(self, key: Union[List, Tuple]) -> None:
+        """Marks a reference as needed.
+
+        Parameters
+        ----------
+        key : Union[List, Tuple]
+            The reference key.
+        """
         assert isinstance(key, (list, tuple)), key
         key = tuple(key)
         self.used_references.add(key)
 
-    def notify_result(self, key, result):
+    def notify_result(self, key: Union[List, Tuple], result: Any) -> None:
+        """Notifies that a result is available for a reference.
+
+        Parameters
+        ----------
+        key : Union[List, Tuple]
+            The reference key.
+        result : Any
+            The result object.
+        """
         trace(
             "🎯",
             step(key),
@@ -45,7 +68,19 @@ class Context:
                 raise ValueError(f"Duplicate result {key}")
             self.results[key] = result
 
-    def get_result(self, key):
+    def get_result(self, key: Union[List, Tuple]) -> Any:
+        """Retrieves the result for a given reference.
+
+        Parameters
+        ----------
+        key : Union[List, Tuple]
+            The reference key.
+
+        Returns
+        -------
+        Any
+            The result for the given reference.
+        """
         assert isinstance(key, (list, tuple)), key
         key = tuple(key)
         if key in self.results:

anemoi/datasets/create/input/data_sources.py

@@ -9,7 +9,14 @@
 
 import logging
 from functools import cached_property
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Union
 
+from earthkit.data import FieldList
+
+from ...dates.groups import GroupOfDates
 from .action import Action
 from .action import action_factory
 from .misc import _tidy
@@ -19,7 +26,28 @@ LOG = logging.getLogger(__name__)
 
 
 class DataSourcesAction(Action):
-    def __init__(self, context, action_path, sources, input):
+    """Class to handle data sources actions in the dataset creation process."""
+
+    def __init__(
+        self,
+        context: object,
+        action_path: List[str],
+        sources: Union[Dict[str, Any], List[Dict[str, Any]]],
+        input: Dict[str, Any],
+    ) -> None:
+        """Initializes a DataSourcesAction instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : List[str]
+            The action path.
+        sources : Union[Dict[str, Any], List[Dict[str, Any]]]
+            The sources configuration.
+        input : Dict[str, Any]
+            The input configuration.
+        """
         super().__init__(context, ["data_sources"], *sources)
         if isinstance(sources, dict):
             configs = [(str(k), c) for k, c in sources.items()]
@@ -31,7 +59,19 @@ class DataSourcesAction(Action):
         self.sources = [action_factory(config, context, ["data_sources"] + [a_path]) for a_path, config in configs]
         self.input = action_factory(input, context, ["input"])
 
-    def select(self, group_of_dates):
+    def select(self, group_of_dates: GroupOfDates) -> "DataSourcesResult":
+        """Selects the data sources result for the given group of dates.
+
+        Parameters
+        ----------
+        group_of_dates : GroupOfDates
+            The group of dates.
+
+        Returns
+        -------
+        DataSourcesResult
+            The data sources result.
+        """
         sources_results = [a.select(group_of_dates) for a in self.sources]
         return DataSourcesResult(
             self.context,
@@ -41,13 +81,38 @@ class DataSourcesAction(Action):
             sources_results,
         )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Returns a string representation of the DataSourcesAction instance."""
         content = "\n".join([str(i) for i in self.sources])
-        return super().__repr__(content)
+        return self._repr(content)
 
 
 class DataSourcesResult(Result):
-    def __init__(self, context, action_path, dates, input_result, sources_results):
+    """Class to represent the result of data sources actions in the dataset creation process."""
+
+    def __init__(
+        self,
+        context: object,
+        action_path: List[str],
+        dates: object,
+        input_result: Result,
+        sources_results: List[Result],
+    ) -> None:
+        """Initializes a DataSourcesResult instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : List[str]
+            The action path.
+        dates : object
+            The dates object.
+        input_result : Result
+            The input result.
+        sources_results : List[Result]
+            The list of sources results.
+        """
         super().__init__(context, action_path, dates)
         # result is the main input result
         self.input_result = input_result
@@ -55,7 +120,8 @@ class DataSourcesResult(Result):
         self.sources_results = sources_results
 
     @cached_property
-    def datasource(self):
+    def datasource(self) -> FieldList:
+        """Returns the combined datasource from all sources."""
         for i in self.sources_results:
             # for each result trigger the datasource to be computed
             # and saved in context

anemoi/datasets/create/input/empty.py

@@ -9,6 +9,9 @@
 
 import logging
 from functools import cached_property
+from typing import List
+
+from earthkit.data import FieldList
 
 from .misc import assert_fieldlist
 from .result import Result
@@ -18,19 +21,34 @@ LOG = logging.getLogger(__name__)
 
 
 class EmptyResult(Result):
+    """Class to represent an empty result in the dataset creation process."""
+
     empty = True
 
-    def __init__(self, context, action_path, dates):
+    def __init__(self, context: object, action_path: list, dates: object) -> None:
+        """Initializes an EmptyResult instance.
+
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : list
+            The action path.
+        dates : object
+            The dates object.
+        """
         super().__init__(context, action_path + ["empty"], dates)
 
     @cached_property
     @assert_fieldlist
     @trace_datasource
-    def datasource(self):
+    def datasource(self) -> FieldList:
+        """Returns an empty datasource."""
         from earthkit.data import from_source
 
         return from_source("empty")
 
     @property
-    def variables(self):
+    def variables(self) -> List[str]:
+        """Returns an empty list of variables."""
         return []

anemoi/datasets/create/input/filter.py

@@ -9,8 +9,11 @@
 
 import logging
 from functools import cached_property
+from typing import Any
+from typing import Type
+
+from earthkit.data import FieldList
 
-from ..functions import import_function
 from .function import FunctionContext
 from .misc import _tidy
 from .misc import assert_fieldlist
@@ -27,14 +30,17 @@ class FilterStepResult(StepResult):
     @notify_result
     @assert_fieldlist
     @trace_datasource
-    def datasource(self):
-        ds = self.upstream_result.datasource
+    def datasource(self) -> FieldList:
+        """Returns the filtered datasource."""
+        ds: FieldList = self.upstream_result.datasource
         ds = ds.sel(**self.action.kwargs)
         return _tidy(ds)
 
 
 class FilterStepAction(StepAction):
-    result_class = FilterStepResult
+    """Represents an action to filter a step result."""
+
+    result_class: Type[FilterStepResult] = FilterStepResult
 
 
 class StepFunctionResult(StepResult):
@@ -42,11 +48,13 @@ class StepFunctionResult(StepResult):
     @assert_fieldlist
     @notify_result
     @trace_datasource
-    def datasource(self):
+    def datasource(self) -> FieldList:
+        """Returns the datasource after applying the function."""
+
+        self.action.filter.context = FunctionContext(self)
         try:
             return _tidy(
-                self.action.function(
-                    FunctionContext(self),
+                self.action.filter.execute(
                     self.upstream_result.datasource,
                     *self.action.args[1:],
                     **self.action.kwargs,
@@ -57,14 +65,54 @@ class StepFunctionResult(StepResult):
             LOG.error(f"Error in {self.action.name}", exc_info=True)
             raise
 
-    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
+            A string representation of the traced datasource.
+        """
         return f"{self.action.name}({self.group_of_dates})"
 
 
 class FunctionStepAction(StepAction):
-    result_class = StepFunctionResult
+    """Represents an action to apply a function to a step result."""
+
+    result_class: Type[StepFunctionResult] = StepFunctionResult
+
+    def __init__(
+        self,
+        context: object,
+        action_path: list,
+        previous_step: StepAction,
+        name: str,
+        filter: Any,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes a FunctionStepAction instance.
 
-    def __init__(self, context, action_path, previous_step, *args, **kwargs):
+        Parameters
+        ----------
+        context : object
+            The context object.
+        action_path : list
+            The action path.
+        previous_step : StepAction
+            The previous step action.
+        *args : Any
+            Additional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
+        """
         super().__init__(context, action_path, previous_step, *args, **kwargs)
-        self.name = args[0]
-        self.function = import_function(self.name, "filters")
+        self.name = name
+        self.filter = filter