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

anemoi/datasets/commands/load-additions.py

@@ -9,6 +9,7 @@
 
 import logging
 import time
+from typing import Any
 
 from anemoi.utils.humanize import seconds_to_human
 
@@ -25,8 +26,14 @@ class LoadAdditions(Command):
     internal = True
     timestamp = True
 
-    def add_arguments(self, subparser):
+    def add_arguments(self, subparser: Any) -> None:
+        """Add command line arguments to the parser.
 
+        Parameters
+        ----------
+        subparser : Any
+            The argument parser.
+        """
         subparser.add_argument("--parts", nargs="+", help="Only load the specified parts of the dataset.")
         subparser.add_argument(
             "--delta",
@@ -37,7 +44,14 @@ class LoadAdditions(Command):
         subparser.add_argument("path", help="Path to store the created data.")
         subparser.add_argument("--trace", action="store_true")
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Execute the command with the given arguments.
+
+        Parameters
+        ----------
+        args : Any
+            The command line arguments.
+        """
         options = vars(args)
         options.pop("command")
         now = time.time()

anemoi/datasets/commands/load.py

@@ -9,6 +9,7 @@
 
 import logging
 import time
+from typing import Any
 
 from anemoi.utils.humanize import seconds_to_human
 
@@ -25,8 +26,14 @@ class Load(Command):
     internal = True
     timestamp = True
 
-    def add_arguments(self, subparser):
+    def add_arguments(self, subparser: Any) -> None:
+        """Add arguments to the command parser.
 
+        Parameters
+        ----------
+        subparser : Any
+            The command parser.
+        """
         subparser.add_argument("--parts", nargs="+", help="Only load the specified parts of the dataset.")
         # subparser.add_argument(
         #        "--delta",
@@ -37,7 +44,14 @@ class Load(Command):
         subparser.add_argument("--cache", help="Location to store the downloaded data.", metavar="DIR")
         subparser.add_argument("--trace", action="store_true")
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Run the command.
+
+        Parameters
+        ----------
+        args : Any
+            The command arguments.
+        """
         options = vars(args)
         options.pop("command")
         now = time.time()

anemoi/datasets/commands/patch.py

@@ -9,6 +9,7 @@
 
 import logging
 import time
+from typing import Any
 
 from anemoi.utils.humanize import seconds_to_human
 
@@ -25,10 +26,24 @@ class Patch(Command):
     internal = True
     timestamp = True
 
-    def add_arguments(self, parser):
+    def add_arguments(self, parser: Any) -> None:
+        """Add command-line arguments to the parser.
+
+        Parameters
+        ----------
+        parser : Any
+            The argument parser instance.
+        """
         parser.add_argument("path", help="Path to store the created data.")
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Execute the patch command.
+
+        Parameters
+        ----------
+        args : Any
+            The command-line arguments.
+        """
         options = vars(args)
         options.pop("command")
         now = time.time()

anemoi/datasets/commands/publish.py

@@ -8,6 +8,7 @@
 # nor does it submit to any jurisdiction.
 
 import logging
+from typing import Any
 
 from . import Command
 
@@ -23,10 +24,24 @@ class Publish(Command):
     internal = True
     timestamp = True
 
-    def add_arguments(self, parser):
+    def add_arguments(self, parser: Any) -> None:
+        """Add arguments to the command parser.
+
+        Parameters
+        ----------
+        parser : Any
+            The command parser to which arguments are added.
+        """
         parser.add_argument("path", help="Path of the dataset to publish.")
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Execute the publish command.
+
+        Parameters
+        ----------
+        args : Any
+            The arguments passed to the command.
+        """
         try:
             from anemoi.registry import publish_dataset
         except ImportError:

anemoi/datasets/commands/scan.py

@@ -11,6 +11,7 @@ import fnmatch
 import os
 import sys
 from collections import defaultdict
+from typing import Any
 
 import earthkit.data as ekd
 import tqdm
@@ -25,7 +26,14 @@ class Scan(Command):
     internal = True
     timestamp = True
 
-    def add_arguments(self, command_parser):
+    def add_arguments(self, command_parser: Any) -> None:
+        """Add arguments to the command parser.
+
+        Parameters
+        ----------
+        command_parser : Any
+            The command parser to which arguments are added.
+        """
         command_parser.add_argument(
             "--match",
             help="Give a glob pattern to match files (default: *.grib)",
@@ -37,8 +45,28 @@ class Scan(Command):
         )
         command_parser.add_argument("paths", nargs="+", help="Paths to scan")
 
-    def run(self, args):
-        def match(path):
+    def run(self, args: Any) -> None:
+        """Execute the scan command.
+
+        Parameters
+        ----------
+        args : Any
+            The arguments passed to the command.
+        """
+
+        def match(path: str) -> bool:
+            """Check if a path matches the given pattern.
+
+            Parameters
+            ----------
+            path : str
+                The path to check.
+
+            Returns
+            -------
+            bool
+                True if the path matches, False otherwise.
+            """
             return fnmatch.fnmatch(path, args.match)
 
         paths = []

anemoi/datasets/compute/recentre.py

@@ -9,13 +9,15 @@
 
 
 import logging
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Tuple
 
 import numpy as np
 from earthkit.data.core.temporary import temp_file
 from earthkit.data.readers.grib.output import new_grib_output
 
-from anemoi.datasets.create.functions import assert_is_fieldlist
-
 LOG = logging.getLogger(__name__)
 
 CLIP_VARIABLES = (
@@ -33,7 +35,22 @@ CLIP_VARIABLES = (
 SKIP = ("class", "stream", "type", "number", "expver", "_leg_number", "anoffset", "time", "date", "step")
 
 
-def check_compatible(f1, f2, centre_field_as_mars, ensemble_field_as_mars):
+def check_compatible(
+    f1: Any, f2: Any, centre_field_as_mars: Dict[str, Any], ensemble_field_as_mars: Dict[str, Any]
+) -> None:
+    """Check if two fields are compatible.
+
+    Parameters
+    ----------
+    f1 : Any
+        The first field.
+    f2 : Any
+        The second field.
+    centre_field_as_mars : Dict[str, Any]
+        Metadata of the centre field.
+    ensemble_field_as_mars : Dict[str, Any]
+        Metadata of the ensemble field.
+    """
     assert f1.mars_grid == f2.mars_grid, (f1.mars_grid, f2.mars_grid)
     assert f1.mars_area == f2.mars_area, (f1.mars_area, f2.mars_area)
     assert f1.shape == f2.shape, (f1.shape, f2.shape)
@@ -56,13 +73,32 @@ def check_compatible(f1, f2, centre_field_as_mars, ensemble_field_as_mars):
 
 def recentre(
     *,
-    members,
-    centre,
-    clip_variables=CLIP_VARIABLES,
-    alpha=1.0,
-    output=None,
-):
-
+    members: Any,
+    centre: Any,
+    clip_variables: Tuple[str, ...] = CLIP_VARIABLES,
+    alpha: float = 1.0,
+    output: Optional[str] = None,
+) -> Any:
+    """Recentre ensemble members around the centre field.
+
+    Parameters
+    ----------
+    members : Any
+        The ensemble members.
+    centre : Any
+        The centre field.
+    clip_variables : Tuple[str, ...], optional
+        Variables to clip. Defaults to CLIP_VARIABLES.
+    alpha : float, optional
+        Scaling factor. Defaults to 1.0.
+    output : Optional[str], optional
+        Output path. Defaults to None.
+
+    Returns
+    -------
+    Any
+        The recentred dataset or output path.
+    """
     keys = ["param", "level", "valid_datetime", "date", "time", "step", "number"]
 
     number_list = members.unique_values("number", progress_bar=False)["number"]
@@ -158,7 +194,7 @@ def recentre(
     from earthkit.data import from_source
 
     ds = from_source("file", path)
-    assert_is_fieldlist(ds)
+
     # save a reference to the tmp file so it is deleted
     # only when the dataset is not used anymore
     ds._tmp = tmp