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

anemoi/datasets/__init__.py

@@ -7,6 +7,8 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+from typing import List
+
 from .data import MissingDateError
 from .data import add_dataset_path
 from .data import add_named_dataset
@@ -21,10 +23,11 @@ except ImportError:  # pragma: no cover
     # Local copy or not installed with setuptools
     __version__ = "999"
 
-__all__ = [
+__all__: List[str] = [
     "add_dataset_path",
     "add_named_dataset",
     "list_dataset_names",
     "MissingDateError",
     "open_dataset",
+    "__version__",
 ]

anemoi/datasets/__main__.py

@@ -7,6 +7,8 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+from typing import Any
+
 from anemoi.utils.cli import cli_main
 from anemoi.utils.cli import make_parser
 
@@ -15,11 +17,19 @@ from .commands import COMMANDS
 
 
 # For read-the-docs
-def create_parser():
+def create_parser() -> Any:
+    """Create the argument parser for the CLI.
+
+    Returns
+    -------
+    Any
+        The argument parser instance.
+    """
     return make_parser(__doc__, COMMANDS)
 
 
-def main():
+def main() -> None:
+    """The main entry point for the CLI application."""
     cli_main(__version__, __doc__, COMMANDS)
 
 

anemoi/datasets/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.16'
-__version_tuple__ = version_tuple = (0, 5, 16)
+__version__ = version = '0.5.18'
+__version_tuple__ = version_tuple = (0, 5, 18)

anemoi/datasets/commands/cleanup.py

@@ -9,6 +9,7 @@
 
 import logging
 import time
+from typing import Any
 
 from anemoi.utils.humanize import seconds_to_human
 
@@ -25,7 +26,14 @@ class Cleanup(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("path", help="Path to store the created data.")
         subparser.add_argument(
             "--delta",
@@ -33,7 +41,14 @@ class Cleanup(Command):
             nargs="+",
         )
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Execute the cleanup command.
+
+        Parameters
+        ----------
+        args : Any
+            The command line arguments.
+        """
         options = vars(args)
         options.pop("command")
         now = time.time()

anemoi/datasets/commands/compare.py

@@ -8,6 +8,8 @@
 # nor does it submit to any jurisdiction.
 
 
+from typing import Any
+
 import numpy as np
 import tqdm
 import zarr
@@ -20,13 +22,27 @@ from . import Command
 class Compare(Command):
     """Compare two datasets. This command compares the variables in two datasets and prints the mean of the common variables. It does not compare the data itself (yet)."""
 
-    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("dataset1")
         command_parser.add_argument("dataset2")
         command_parser.add_argument("--data", action="store_true", help="Compare the data.")
         command_parser.add_argument("--statistics", action="store_true", help="Compare the statistics.")
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Run the compare command with the provided arguments.
+
+        Parameters
+        ----------
+        args : Any
+            The arguments passed to the command.
+        """
         ds1 = open_dataset(args.dataset1)
         ds2 = open_dataset(args.dataset2)
 

anemoi/datasets/commands/copy.py

@@ -13,6 +13,8 @@ import os
 import sys
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import as_completed
+from typing import Any
+from typing import Optional
 
 import tqdm
 from anemoi.utils.remote import Transfer
@@ -29,7 +31,68 @@ except AttributeError:
 
 
 class ZarrCopier:
-    def __init__(self, source, target, transfers, block_size, overwrite, resume, verbosity, nested, rechunk, **kwargs):
+    """Class to handle copying of Zarr datasets.
+
+    Attributes
+    ----------
+    source : str
+        Source location of the dataset.
+    target : str
+        Target location of the dataset.
+    transfers : int
+        Number of parallel transfers.
+    block_size : int
+        Size of data blocks to transfer.
+    overwrite : bool
+        Flag to overwrite existing dataset.
+    resume : bool
+        Flag to resume copying an existing dataset.
+    verbosity : int
+        Verbosity level of logging.
+    nested : bool
+        Flag to use ZARR's nested directory backend.
+    rechunk : str
+        Rechunk size for the target data array.
+    """
+
+    def __init__(
+        self,
+        source: str,
+        target: str,
+        transfers: int,
+        block_size: int,
+        overwrite: bool,
+        resume: bool,
+        verbosity: int,
+        nested: bool,
+        rechunk: str,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the ZarrCopier.
+
+        Parameters
+        ----------
+        source : str
+            Source location of the dataset.
+        target : str
+            Target location of the dataset.
+        transfers : int
+            Number of parallel transfers.
+        block_size : int
+            Size of data blocks to transfer.
+        overwrite : bool
+            Flag to overwrite existing dataset.
+        resume : bool
+            Flag to resume copying an existing dataset.
+        verbosity : int
+            Verbosity level of logging.
+        nested : bool
+            Flag to use ZARR's nested directory backend.
+        rechunk : str
+            Rechunk size for the target data array.
+        **kwargs : Any
+            Additional keyword arguments.
+        """
         self.source = source
         self.target = target
         self.transfers = transfers
@@ -50,14 +113,50 @@ class ZarrCopier:
                 raise NotImplementedError("Rechunking with SSH not implemented.")
             assert NotImplementedError("SSH not implemented.")
 
-    def _store(self, path, nested=False):
+    def _store(self, path: str, nested: bool = False) -> Any:
+        """Get the storage path.
+
+        Parameters
+        ----------
+        path : str
+            Path to the storage.
+        nested : bool, optional
+            Flag to use nested directory storage.
+
+        Returns
+        -------
+        Any
+            Storage path.
+        """
         if nested:
             import zarr
 
             return zarr.storage.NestedDirectoryStore(path)
         return path
 
-    def copy_chunk(self, n, m, source, target, _copy, verbosity):
+    def copy_chunk(self, n: int, m: int, source: Any, target: Any, _copy: Any, verbosity: int) -> Optional[slice]:
+        """Copy a chunk of data from source to target.
+
+        Parameters
+        ----------
+        n : int
+            Start index of the chunk.
+        m : int
+            End index of the chunk.
+        source : Any
+            Source data.
+        target : Any
+            Target data.
+        _copy : Any
+            Copy status array.
+        verbosity : int
+            Verbosity level of logging.
+
+        Returns
+        -------
+        slice or None
+            Slice of copied data or None if skipped.
+        """
         if _copy[n:m].all():
             LOG.info(f"Skipping {n} to {m}")
             return None
@@ -83,7 +182,21 @@ class ZarrCopier:
 
         return slice(n, m)
 
-    def parse_rechunking(self, rechunking, source_data):
+    def parse_rechunking(self, rechunking: list[str], source_data: Any) -> tuple:
+        """Parse the rechunking configuration.
+
+        Parameters
+        ----------
+        rechunking : list of str
+            List of rechunk sizes.
+        source_data : Any
+            Source data.
+
+        Returns
+        -------
+        tuple
+            Parsed chunk sizes.
+        """
         shape = source_data.shape
         chunks = list(source_data.chunks)
         for i, c in enumerate(rechunking):
@@ -102,7 +215,20 @@ class ZarrCopier:
             #    raise NotImplementedError("Rechunking with multiple transfers is not implemented")
         return chunks
 
-    def copy_data(self, source, target, _copy, verbosity):
+    def copy_data(self, source: Any, target: Any, _copy: Any, verbosity: int) -> None:
+        """Copy data from source to target.
+
+        Parameters
+        ----------
+        source : Any
+            Source data.
+        target : Any
+            Target data.
+        _copy : Any
+            Copy status array.
+        verbosity : int
+            Verbosity level of logging.
+        """
         LOG.info("Copying data")
         source_data = source["data"]
 
@@ -147,7 +273,22 @@ class ZarrCopier:
 
         LOG.info("Copied data")
 
-    def copy_array(self, name, source, target, _copy, verbosity):
+    def copy_array(self, name: str, source: Any, target: Any, _copy: Any, verbosity: int) -> None:
+        """Copy an array from source to target.
+
+        Parameters
+        ----------
+        name : str
+            Name of the array.
+        source : Any
+            Source data.
+        target : Any
+            Target data.
+        _copy : Any
+            Copy status array.
+        verbosity : int
+            Verbosity level of logging.
+        """
         for k, v in source.attrs.items():
             target.attrs[k] = v
 
@@ -162,7 +303,20 @@ class ZarrCopier:
         target[name] = source[name]
         LOG.info(f"Copied {name}")
 
-    def copy_group(self, source, target, _copy, verbosity):
+    def copy_group(self, source: Any, target: Any, _copy: Any, verbosity: int) -> None:
+        """Copy a group from source to target.
+
+        Parameters
+        ----------
+        source : Any
+            Source data.
+        target : Any
+            Target data.
+        _copy : Any
+            Copy status array.
+        verbosity : int
+            Verbosity level of logging.
+        """
         import zarr
 
         for k, v in source.attrs.items():
@@ -186,7 +340,18 @@ class ZarrCopier:
                     verbosity,
                 )
 
-    def copy(self, source, target, verbosity):
+    def copy(self, source: Any, target: Any, verbosity: int) -> None:
+        """Copy the entire dataset from source to target.
+
+        Parameters
+        ----------
+        source : Any
+            Source data.
+        target : Any
+            Target data.
+        verbosity : int
+            Verbosity level of logging.
+        """
         import zarr
 
         if "_copy" not in target:
@@ -200,7 +365,8 @@ class ZarrCopier:
         self.copy_group(source, target, _copy_np, verbosity)
         del target["_copy"]
 
-    def run(self):
+    def run(self) -> None:
+        """Execute the copy operation."""
         import zarr
 
         # base, ext = os.path.splitext(os.path.basename(args.source))
@@ -208,14 +374,14 @@ class ZarrCopier:
         # assert "." not in base, base
         LOG.info(f"Copying {self.source} to {self.target}")
 
-        def target_exists():
+        def target_exists() -> bool:
             try:
                 zarr.open(self._store(self.target), mode="r")
                 return True
             except ValueError:
                 return False
 
-        def target_finished():
+        def target_finished() -> bool:
             target = zarr.open(self._store(self.target), mode="r")
             if "_copy" in target:
                 done = sum(1 if x else 0 for x in target["_copy"])
@@ -231,7 +397,7 @@ class ZarrCopier:
                 return True
             return False
 
-        def open_target():
+        def open_target() -> Any:
 
             if not target_exists():
                 return zarr.open(self._store(self.target, self.nested), mode="w")
@@ -260,10 +426,19 @@ class ZarrCopier:
 
 
 class CopyMixin:
+    """Mixin class for adding copy command arguments and running the copy operation."""
+
     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
+            Command parser object.
+        """
         group = command_parser.add_mutually_exclusive_group()
         group.add_argument(
             "--overwrite",
@@ -293,7 +468,14 @@ class CopyMixin:
         command_parser.add_argument("source", help="Source location.")
         command_parser.add_argument("target", help="Target location.")
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Run the copy command with the provided arguments.
+
+        Parameters
+        ----------
+        args : Any
+            Command arguments.
+        """
         if args.source == args.target:
             raise ValueError("Source and target are the same.")
 

anemoi/datasets/commands/create.py

@@ -13,6 +13,7 @@ import time
 from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import as_completed
+from typing import Any
 
 import tqdm
 from anemoi.utils.humanize import seconds_to_human
@@ -22,9 +23,25 @@ from . import Command
 LOG = logging.getLogger(__name__)
 
 
-def task(what, options, *args, **kwargs):
-    """Make sure `import Creator` is done in the sub-processes, and not in the main one."""
-
+def task(what: str, options: dict, *args: Any, **kwargs: Any) -> Any:
+    """Make sure `import Creator` is done in the sub-processes, and not in the main one.
+
+    Parameters
+    ----------
+    what : str
+        The task to be executed.
+    options : dict
+        Options for the task.
+    *args : Any
+        Additional arguments.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    Any
+        The result of the task.
+    """
     now = datetime.datetime.now()
     LOG.info(f"🎬 Task {what}({args},{kwargs}) starting")
 
@@ -45,7 +62,14 @@ class Create(Command):
     internal = True
     timestamp = True
 
-    def add_arguments(self, command_parser):
+    def add_arguments(self, command_parser: Any) -> None:
+        """Add command line arguments to the parser.
+
+        Parameters
+        ----------
+        command_parser : Any
+            The command line argument parser.
+        """
         command_parser.add_argument(
             "--overwrite",
             action="store_true",
@@ -63,8 +87,14 @@ class Create(Command):
         group.add_argument("--processes", help="Use `n` parallel process workers.", type=int, default=0)
         command_parser.add_argument("--trace", action="store_true")
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Execute the create command.
 
+        Parameters
+        ----------
+        args : Any
+            Command line arguments.
+        """
         now = time.time()
         if args.threads + args.processes:
             self.parallel_create(args)
@@ -72,8 +102,14 @@ class Create(Command):
             self.serial_create(args)
         LOG.info(f"Create completed in {seconds_to_human(time.time()-now)}")
 
-    def serial_create(self, args):
+    def serial_create(self, args: Any) -> None:
+        """Create the dataset in serial mode.
 
+        Parameters
+        ----------
+        args : Any
+            Command line arguments.
+        """
         options = vars(args)
         options.pop("command")
         options.pop("threads")
@@ -92,7 +128,14 @@ class Create(Command):
         task("cleanup", options)
         task("verify", options)
 
-    def parallel_create(self, args):
+    def parallel_create(self, args: Any) -> None:
+        """Create the dataset in parallel mode.
+
+        Parameters
+        ----------
+        args : Any
+            Command line arguments.
+        """
         """Some modules, like fsspec do not work well with fork()
         Other modules may not be thread safe. So we implement
         parallel loadining using multiprocessing before any

anemoi/datasets/commands/finalise-additions.py

@@ -9,6 +9,7 @@
 
 import logging
 import time
+from typing import Any
 
 from anemoi.utils.humanize import seconds_to_human
 
@@ -25,7 +26,14 @@ class FinaliseAdditions(Command):
     internal = True
     timestamp = True
 
-    def add_arguments(self, command_parser):
+    def add_arguments(self, command_parser: Any) -> None:
+        """Add command line arguments to the parser.
+
+        Parameters
+        ----------
+        command_parser : Any
+            The argument parser instance to which arguments will be added.
+        """
         command_parser.add_argument(
             "--delta",
             help="Compute statistics tendencies on a given time delta, if possible. Must be a multiple of the frequency.",
@@ -35,7 +43,14 @@ class FinaliseAdditions(Command):
         command_parser.add_argument("path", help="Path to store the created data.")
         command_parser.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 arguments passed to the command.
+        """
         options = vars(args)
         options.pop("command")
         step = "finalise-additions"

anemoi/datasets/commands/finalise.py

@@ -9,6 +9,7 @@
 
 import logging
 import time
+from typing import Any
 
 from anemoi.utils.humanize import seconds_to_human
 
@@ -25,11 +26,25 @@ class Finalise(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 will be added.
+        """
         command_parser.add_argument("path", help="Path to store the created data.")
         command_parser.add_argument("--trace", action="store_true")
 
-    def run(self, args):
+    def run(self, args: Any) -> None:
+        """Execute the finalise command.
+
+        Parameters
+        ----------
+        args : Any
+            The arguments passed to the command.
+        """
         options = vars(args)
         options.pop("command")
         now = time.time()

anemoi/datasets/commands/init-additions.py

@@ -9,6 +9,7 @@
 
 import logging
 import time
+from typing import Any
 
 from anemoi.utils.humanize import seconds_to_human
 
@@ -25,7 +26,14 @@ class InitAdditions(Command):
     internal = True
     timestamp = True
 
-    def add_arguments(self, command_parser):
+    def add_arguments(self, command_parser: Any) -> None:
+        """Add command line arguments to the parser.
+
+        Parameters
+        ----------
+        command_parser : Any
+            The argument parser instance.
+        """
         command_parser.add_argument(
             "--delta",
             help="Compute statistics tendencies on a given time delta, if possible. Must be a multiple of the frequency.",
@@ -35,7 +43,14 @@ class InitAdditions(Command):
         command_parser.add_argument("path", help="Path to store the created data.")
         command_parser.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")
         step = "init-additions"