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

anemoi/datasets/create/input/step.py

@@ -8,11 +8,17 @@
 # nor does it submit to any jurisdiction.
 
 import logging
+import warnings
 from copy import deepcopy
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Type
 
 from .action import Action
+from .action import ActionContext
 from .context import Context
-from .misc import is_function
 from .result import Result
 from .template import notify_result
 from .trace import trace_datasource
@@ -22,28 +28,75 @@ LOG = logging.getLogger(__name__)
 
 
 class StepResult(Result):
-    def __init__(self, context, action_path, group_of_dates, action, upstream_result):
+    """Represents the result of a step in the data processing pipeline."""
+
+    def __init__(
+        self, context: Context, action_path: List[str], group_of_dates: Any, action: Action, upstream_result: Result
+    ) -> None:
+        """Initialize a StepResult instance.
+
+        Parameters
+        ----------
+        context
+            The context in which the step is executed.
+        action_path
+            The path of actions leading to this step.
+        group_of_dates
+            The group of dates associated with this step.
+        action
+            The action associated with this step.
+        upstream_result
+            The result of the upstream step.
+        """
         super().__init__(context, action_path, group_of_dates)
         assert isinstance(upstream_result, Result), type(upstream_result)
-        self.upstream_result = upstream_result
-        self.action = action
+        self.upstream_result: Result = upstream_result
+        self.action: Action = action
 
     @property
     @notify_result
     @trace_datasource
-    def datasource(self):
+    def datasource(self) -> Any:
+        """Retrieve the datasource associated with this step result."""
         raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
 
 
 class StepAction(Action):
-    result_class = None
-
-    def __init__(self, context, action_path, previous_step, *args, **kwargs):
+    """Represents an action that is part of a step in the data processing pipeline."""
+
+    result_class: Optional[Type[StepResult]] = None
+
+    def __init__(
+        self, context: ActionContext, action_path: List[str], previous_step: Any, *args: Any, **kwargs: Any
+    ) -> None:
+        """Initialize a StepAction instance.
+
+        Parameters
+        ----------
+        context
+            The context in which the action is executed.
+        action_path
+            The path of actions leading to this step.
+        previous_step
+            The previous step in the pipeline.
+        """
         super().__init__(context, action_path, *args, **kwargs)
-        self.previous_step = previous_step
+        self.previous_step: Any = previous_step
 
     @trace_select
-    def select(self, group_of_dates):
+    def select(self, group_of_dates: Any) -> StepResult:
+        """Select the result for a given group of dates.
+
+        Parameters
+        ----------
+        group_of_dates
+            The group of dates to select the result for.
+
+        Returns
+        -------
+        unknown
+            The result of the step.
+        """
         return self.result_class(
             self.context,
             self.action_path,
@@ -52,12 +105,36 @@ class StepAction(Action):
             self.previous_step.select(group_of_dates),
         )
 
-    def __repr__(self):
-        # raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
-        return super().__repr__(self.previous_step, _inline_=str(self.kwargs))
-
-
-def step_factory(config, context, action_path, previous_step):
+    def __repr__(self) -> str:
+        """Return a string representation of the StepAction instance.
+
+        Returns
+        -------
+        unknown
+            String representation of the instance.
+        """
+        return self._repr(self.previous_step, _inline_=str(self.kwargs))
+
+
+def step_factory(config: Dict[str, Any], context: ActionContext, action_path: List[str], previous_step: Any) -> Any:
+    """Factory function to create a step action based on the given configuration.
+
+    Parameters
+    ----------
+    config
+        The configuration dictionary for the step.
+    context
+        The context in which the step is executed.
+    action_path
+        The path of actions leading to this step.
+    previous_step
+        The previous step in the pipeline.
+
+    Returns
+    -------
+    unknown
+        An instance of a step action.
+    """
 
     from .filter import FilterStepAction
     from .filter import FunctionStepAction
@@ -85,10 +162,30 @@ def step_factory(config, context, action_path, previous_step):
     if isinstance(config[key], str):
         args, kwargs = [config[key]], {}
 
-    if cls is None:
-        if not is_function(key, "filters"):
-            raise ValueError(f"Unknown step {key}")
-        cls = FunctionStepAction
-        args = [key] + args
+    if cls is not None:
+        return cls(context, action_path, previous_step, *args, **kwargs)
+
+    # Try filters from datasets filter registry
+    from anemoi.transform.filters import filter_registry as transform_filter_registry
+
+    from ..filters import create_filter as create_datasets_filter
+    from ..filters import filter_registry as datasets_filter_registry
+
+    if datasets_filter_registry.is_registered(key):
+
+        if transform_filter_registry.is_registered(key):
+            warnings.warn(f"Filter `{key}` is registered in both datasets and transform filter registries")
+
+        filter = create_datasets_filter(None, config)
+        return FunctionStepAction(context, action_path + [key], previous_step, key, filter)
+
+    # Use filters from transform registry
+
+    if transform_filter_registry.is_registered(key):
+        from ..filters.transform import TransformFilter
+
+        return FunctionStepAction(
+            context, action_path + [key], previous_step, key, TransformFilter(context, key, config)
+        )
 
-    return cls(context, action_path, previous_step, *args, **kwargs)
+    raise ValueError(f"Unknown step action `{key}`")

anemoi/datasets/create/input/template.py

@@ -10,35 +10,108 @@
 
 import logging
 import re
+from abc import ABC
+from abc import abstractmethod
 from functools import wraps
+from typing import Any
+from typing import Callable
+from typing import List
+
+from .context import Context
 
 LOG = logging.getLogger(__name__)
 
 
-def notify_result(method):
+def notify_result(method: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator to notify the context of the result of the method call.
+
+    Parameters
+    ----------
+    method : Callable[..., Any]
+        The method to wrap.
+
+    Returns
+    -------
+    Callable[..., Any]
+        The wrapped method.
+    """
+
     @wraps(method)
-    def wrapper(self, *args, **kwargs):
-        result = method(self, *args, **kwargs)
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        result: Any = method(self, *args, **kwargs)
         self.context.notify_result(self.action_path, result)
         return result
 
     return wrapper
 
 
-class Substitution:
-    pass
+class Substitution(ABC):
+    """Abstract base class for substitutions in templates."""
 
+    @abstractmethod
+    def resolve(self, context: Context) -> Any:
+        """Resolve the substitution using the given context.
 
-class Reference(Substitution):
-    def __init__(self, context, action_path):
-        self.context = context
-        self.action_path = action_path
+        Parameters
+        ----------
+        context : Context
+            The context to use for resolution.
+
+        Returns
+        -------
+        Any
+            The resolved value.
+        """
+        pass
 
-    def resolve(self, context):
+
+class Reference(Substitution):
+    """A class to represent a reference to another value in the context."""
+
+    def __init__(self, context: Any, action_path: List[str]) -> None:
+        """Initialize a Reference instance.
+
+        Parameters
+        ----------
+        context : Any
+            The context in which the reference exists.
+        action_path : list of str
+            The action path to resolve.
+        """
+        self.context: Any = context
+        self.action_path: List[str] = action_path
+
+    def resolve(self, context: Context) -> Any:
+        """Resolve the reference using the given context.
+
+        Parameters
+        ----------
+        context : Context
+            The context to use for resolution.
+
+        Returns
+        -------
+        Any
+            The resolved value.
+        """
         return context.get_result(self.action_path)
 
 
-def resolve(context, x):
+def resolve(context: Context, x: Any) -> Any:
+    """Recursively resolve substitutions in the given structure using the context.
+
+    Parameters
+    ----------
+    context : Context
+        The context to use for resolution.
+    x : Union[tuple, list, dict, Substitution, Any]
+        The structure to resolve.
+
+    Returns
+    -------
+    Any
+        The resolved structure.
+    """
     if isinstance(x, tuple):
         return tuple([resolve(context, y) for y in x])
 
@@ -54,7 +127,21 @@ def resolve(context, x):
     return x
 
 
-def substitute(context, x):
+def substitute(context: Context, x: Any) -> Any:
+    """Recursively substitute references in the given structure using the context.
+
+    Parameters
+    ----------
+    context : Context
+        The context to use for substitution.
+    x : Union[tuple, list, dict, str, Any]
+        The structure to substitute.
+
+    Returns
+    -------
+    Any
+        The substituted structure.
+    """
     if isinstance(x, tuple):
         return tuple([substitute(context, y) for y in x])
 
@@ -67,7 +154,7 @@ def substitute(context, x):
     if not isinstance(x, str):
         return x
 
-    if re.match(r"^\${[\.\w]+}$", x):
+    if re.match(r"^\${[\.\w\-]+}$", x):
         path = x[2:-1].split(".")
         context.will_need_reference(path)
         return Reference(context, path)

anemoi/datasets/create/input/trace.py

@@ -12,6 +12,8 @@ import logging
 import textwrap
 import threading
 from functools import wraps
+from typing import Any
+from typing import Callable
 
 LOG = logging.getLogger(__name__)
 
@@ -20,17 +22,44 @@ thread_local = threading.local()
 TRACE = 0
 
 
-def enable_trace(on_off):
+def enable_trace(on_off: int) -> None:
+    """Enables or disables tracing.
+
+    Parameters
+    ----------
+    on_off : int
+        1 to enable tracing, 0 to disable.
+    """
     global TRACE
     TRACE = on_off
 
 
-def step(action_path):
+def step(action_path: list[str]) -> str:
+    """Returns a formatted string representing the action path.
+
+    Parameters
+    ----------
+    action_path : list of str
+        The action path.
+
+    Returns
+    -------
+    str
+        The formatted action path.
+    """
     return f"[{'.'.join(action_path)}]"
 
 
-def trace(emoji, *args):
+def trace(emoji: str, *args: Any) -> None:
+    """Logs a trace message with an emoji.
 
+    Parameters
+    ----------
+    emoji : str
+        The emoji to use.
+    *args : Any
+        The arguments to log.
+    """
     if not TRACE:
         return
 
@@ -40,9 +69,22 @@ def trace(emoji, *args):
     print(emoji, " " * thread_local.TRACE_INDENT, *args)
 
 
-def trace_datasource(method):
+def trace_datasource(method: Callable) -> Callable:
+    """Decorator to trace the datasource method.
+
+    Parameters
+    ----------
+    method : Callable
+        The method to decorate.
+
+    Returns
+    -------
+    Callable
+        The wrapped method.
+    """
+
     @wraps(method)
-    def wrapper(self, *args, **kwargs):
+    def wrapper(self, *args: Any, **kwargs: Any) -> Any:
 
         if not hasattr(thread_local, "TRACE_INDENT"):
             thread_local.TRACE_INDENT = 0
@@ -67,9 +109,22 @@ def trace_datasource(method):
     return wrapper
 
 
-def trace_select(method):
+def trace_select(method: Callable) -> Callable:
+    """Decorator to trace the select method.
+
+    Parameters
+    ----------
+    method : Callable
+        The method to decorate.
+
+    Returns
+    -------
+    Callable
+        The wrapped method.
+    """
+
     @wraps(method)
-    def wrapper(self, *args, **kwargs):
+    def wrapper(self, *args: Any, **kwargs: Any) -> Any:
         if not hasattr(thread_local, "TRACE_INDENT"):
             thread_local.TRACE_INDENT = 0
         trace(

anemoi/datasets/create/patch.py

@@ -10,13 +10,26 @@
 import json
 import logging
 import os
+from typing import Union
 
 import zarr
 
 LOG = logging.getLogger(__name__)
 
 
-def fix_order_by(order_by):
+def fix_order_by(order_by: Union[dict, list]) -> list[dict]:
+    """Fix the order_by attribute to ensure it is a list of dictionaries.
+
+    Parameters
+    ----------
+    order_by : dict or list
+        The order_by attribute to fix.
+
+    Returns
+    -------
+    list[dict]
+        The fixed order_by attribute.
+    """
     if isinstance(order_by, list):
         return order_by
 
@@ -29,14 +42,38 @@ def fix_order_by(order_by):
     return lst
 
 
-def fix_history(history):
+def fix_history(history: list[dict]) -> list[dict]:
+    """Fix the history attribute by removing specific actions.
+
+    Parameters
+    ----------
+    history : list[dict]
+        The history attribute to fix.
+
+    Returns
+    -------
+    list[dict]
+        The fixed history attribute.
+    """
     new = history
     new = [d for d in new if d.get("action") != "loading_data_start"]
     new = [d for d in new if d.get("action") != "loading_data_end"]
     return new
 
 
-def fix_provenance(provenance):
+def fix_provenance(provenance: dict) -> dict:
+    """Fix the provenance attribute by adding missing fields and removing unnecessary ones.
+
+    Parameters
+    ----------
+    provenance : dict
+        The provenance attribute to fix.
+
+    Returns
+    -------
+    dict
+        The fixed provenance attribute.
+    """
     if "python" not in provenance:
         provenance["python"] = provenance["platform"]["python_version"]
 
@@ -79,7 +116,18 @@ def fix_provenance(provenance):
     return provenance
 
 
-def apply_patch(path, verbose=True, dry_run=False):
+def apply_patch(path: str, verbose: bool = True, dry_run: bool = False) -> None:
+    """Apply a patch to the dataset at the given path.
+
+    Parameters
+    ----------
+    path : str
+        The path to the dataset.
+    verbose : bool, optional
+        Whether to log detailed information. Defaults to True.
+    dry_run : bool, optional
+        If True, do not actually apply the patch. Defaults to False.
+    """
     LOG.debug("====================")
     LOG.debug(f"Patching {path}")
     LOG.debug("====================")