anemoi-utils 0.4.12__py3-none-any.whl → 0.4.14__py3-none-any.whl

anemoi/utils/remote/__init__.py

@@ -10,6 +10,9 @@ import logging
 import os
 import shutil
 from abc import abstractmethod
+from typing import Any
+from typing import Dict
+from typing import Iterable
 
 import tqdm
 
@@ -18,13 +21,63 @@ from ..humanize import bytes_to_human
 LOGGER = logging.getLogger(__name__)
 
 
-def _ignore(number_of_files, total_size, total_transferred, transfering):
+def robust(call: callable, *args, maximum_tries: int = 60, retry_after: int = 60, **kwargs) -> callable:
+    """Forwards the arguments to the multiurl robust function.
+    with default retry_after=60 and maximum_tries=60.
+    """
+    from multiurl import robust as robust_
+
+    return robust_(call, *args, retry_after=retry_after, maximum_tries=maximum_tries, **kwargs)
+
+
+def _ignore(number_of_files: int, total_size: int, total_transferred: int, transfering: bool) -> None:
+    """A placeholder function for progress reporting.
+
+    Parameters
+    ----------
+    number_of_files : int
+        The number of files being transferred.
+    total_size : int
+        The total size of the files being transferred.
+    total_transferred : int
+        The total size of the files transferred so far.
+    transfering : bool
+        Whether the transfer is in progress.
+    """
     pass
 
 
 class Loader:
-
-    def transfer_folder(self, *, source, target, overwrite=False, resume=False, verbosity=1, threads=1, progress=None):
+    def transfer_folder(
+        self,
+        *,
+        source: str,
+        target: str,
+        overwrite: bool = False,
+        resume: bool = False,
+        verbosity: int = 1,
+        threads: int = 1,
+        progress: callable = None,
+    ) -> None:
+        """Transfer a folder from the source to the target location.
+
+        Parameters
+        ----------
+        source : str
+            The source folder path.
+        target : str
+            The target folder path.
+        overwrite : bool, optional
+            Whether to overwrite the target if it exists, by default False.
+        resume : bool, optional
+            Whether to resume the transfer if possible, by default False.
+        verbosity : int, optional
+            The verbosity level, by default 1.
+        threads : int, optional
+            The number of threads to use, by default 1.
+        progress : callable, optional
+            A callable for progress reporting, by default None.
+        """
         assert verbosity == 1, verbosity
 
         if progress is None:
@@ -93,7 +146,48 @@ class Loader:
                 executor.shutdown(wait=False, cancel_futures=True)
                 raise
 
-    def transfer_file(self, source, target, overwrite, resume, verbosity, threads=1, progress=None, config=None):
+    def transfer_file(
+        self,
+        source: str,
+        target: str,
+        overwrite: bool,
+        resume: bool,
+        verbosity: int,
+        threads: int = 1,
+        progress: callable = None,
+        config: dict = None,
+    ) -> int:
+        """Transfer a file from the source to the target location.
+
+        Parameters
+        ----------
+        source : str
+            The source file path.
+        target : str
+            The target file path.
+        overwrite : bool
+            Whether to overwrite the target if it exists.
+        resume : bool
+            Whether to resume the transfer if possible.
+        verbosity : int
+            The verbosity level.
+        threads : int, optional
+            The number of threads to use, by default 1.
+        progress : callable, optional
+            A callable for progress reporting, by default None.
+        config : dict, optional
+            Additional configuration options, by default None.
+
+        Returns
+        -------
+        int
+            The size of the transferred file.
+
+        Raises
+        ------
+        Exception
+            If an error occurs during the transfer.
+        """
         try:
             return self._transfer_file(source, target, overwrite, resume, verbosity, threads=threads, config=config)
         except Exception as e:
@@ -102,31 +196,119 @@ class Loader:
             raise
 
     @abstractmethod
-    def list_source(self, source):
+    def list_source(self, source: str) -> Iterable:
+        """List the files in the source location.
+
+        Parameters
+        ----------
+        source : str
+            The source location.
+
+        Returns
+        -------
+        Iterable
+            An iterable of files in the source location.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def source_path(self, local_path, source):
+    def source_path(self, local_path: str, source: str) -> str:
+        """Get the source path for a local file.
+
+        Parameters
+        ----------
+        local_path : str
+            The local file path.
+        source : str
+            The source location.
+
+        Returns
+        -------
+        str
+            The source path for the local file.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def target_path(self, source_path, source, target):
+    def target_path(self, source_path: str, source: str, target: str) -> str:
+        """Get the target path for a source file.
+
+        Parameters
+        ----------
+        source_path : str
+            The source file path.
+        source : str
+            The source location.
+        target : str
+            The target location.
+
+        Returns
+        -------
+        str
+            The target path for the source file.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def source_size(self, local_path):
+    def source_size(self, local_path: str) -> int:
+        """Get the size of a local file.
+
+        Parameters
+        ----------
+        local_path : str
+            The local file path.
+
+        Returns
+        -------
+        int
+            The size of the local file.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def copy(self, source, target, **kwargs):
+    def copy(self, source: str, target: str, **kwargs) -> None:
+        """Copy a file or folder from the source to the target location.
+
+        Parameters
+        ----------
+        source : str
+            The source location.
+        target : str
+            The target location.
+        kwargs : dict
+            Additional arguments for the transfer.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def get_temporary_target(self, target, pattern):
+    def get_temporary_target(self, target: str, pattern: str) -> str:
+        """Get a temporary target path based on the given pattern.
+
+        Parameters
+        ----------
+        target : str
+            The original target path.
+        pattern : str
+            The pattern to format the temporary path.
+
+        Returns
+        -------
+        str
+            The temporary target path.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def rename_target(self, target, temporary_target):
+    def rename_target(self, target: str, temporary_target: str) -> None:
+        """Rename the target to a new target path.
+
+        Parameters
+        ----------
+        target : str
+            The original target path.
+        temporary_target : str
+            The new target path.
+        """
         raise NotImplementedError
 
 
@@ -134,19 +316,60 @@ class BaseDownload(Loader):
     action = "Downloading"
 
     @abstractmethod
-    def copy(self, source, target, **kwargs):
+    def copy(self, source: str, target: str, **kwargs) -> None:
+        """Copy a file or folder from the source to the target location.
+
+        Parameters
+        ----------
+        source : str
+            The source location.
+        target : str
+            The target location.
+        kwargs : dict
+            Additional arguments for the transfer.
+        """
         raise NotImplementedError
 
-    def get_temporary_target(self, target, pattern):
+    def get_temporary_target(self, target: str, pattern: str) -> str:
+        """Get a temporary target path based on the given pattern.
+
+        Parameters
+        ----------
+        target : str
+            The original target path.
+        pattern : str
+            The pattern to format the temporary path.
+
+        Returns
+        -------
+        str
+            The temporary target path.
+        """
         if pattern is None:
             return target
         dirname, basename = os.path.split(target)
         return pattern.format(dirname=dirname, basename=basename)
 
-    def rename_target(self, target, new_target):
+    def rename_target(self, target: str, new_target: str) -> None:
+        """Rename the target to a new target path.
+
+        Parameters
+        ----------
+        target : str
+            The original target path.
+        new_target : str
+            The new target path.
+        """
         os.rename(target, new_target)
 
-    def delete_target(self, target):
+    def delete_target(self, target: str) -> None:
+        """Delete the target if it exists.
+
+        Parameters
+        ----------
+        target : str
+            The target path.
+        """
         if os.path.exists(target):
             shutil.rmtree(target)
 
@@ -154,26 +377,91 @@ class BaseDownload(Loader):
 class BaseUpload(Loader):
     action = "Uploading"
 
-    def copy(self, source, target, **kwargs):
+    def copy(self, source: str, target: str, **kwargs) -> None:
+        """Copy a file or folder from the source to the target location.
+
+        Parameters
+        ----------
+        source : str
+            The source location.
+        target : str
+            The target location.
+        kwargs : dict
+            Additional arguments for the transfer.
+        """
         if os.path.isdir(source):
             self.transfer_folder(source=source, target=target, **kwargs)
         else:
             self.transfer_file(source=source, target=target, **kwargs)
 
-    def list_source(self, source):
+    def list_source(self, source: str) -> Iterable:
+        """List the files in the source location.
+
+        Parameters
+        ----------
+        source : str
+            The source location.
+
+        Returns
+        -------
+        Iterable
+            An iterable of files in the source location.
+        """
         for root, _, files in os.walk(source):
             for file in files:
                 yield os.path.join(root, file)
 
-    def source_path(self, local_path, source):
+    def source_path(self, local_path: str, source: str) -> str:
+        """Get the source path for a local file.
+
+        Parameters
+        ----------
+        local_path : str
+            The local file path.
+        source : str
+            The source location.
+
+        Returns
+        -------
+        str
+            The source path for the local file.
+        """
         return local_path
 
-    def target_path(self, source_path, source, target):
+    def target_path(self, source_path: str, source: str, target: str) -> str:
+        """Get the target path for a source file.
+
+        Parameters
+        ----------
+        source_path : str
+            The source file path.
+        source : str
+            The source location.
+        target : str
+            The target location.
+
+        Returns
+        -------
+        str
+            The target path for the source file.
+        """
         relative_path = os.path.relpath(source_path, source)
         path = os.path.join(target, relative_path)
         return path
 
-    def source_size(self, local_path):
+    def source_size(self, local_path: str) -> int:
+        """Get the size of a local file.
+
+        Parameters
+        ----------
+        local_path : str
+            The local file path.
+
+        Returns
+        -------
+        int
+            The size of the local file.
+        """
         return os.path.getsize(local_path)
 
 
@@ -188,21 +476,22 @@ class Transfer:
 
     def __init__(
         self,
-        source,
-        target,
-        overwrite=False,
-        resume=False,
-        verbosity=1,
-        threads=1,
-        progress=None,
-        temporary_target=False,
+        *,
+        source: str,
+        target: str,
+        overwrite: bool = False,
+        resume: bool = False,
+        verbosity: int = 1,
+        threads: int = 1,
+        progress: callable = None,
+        temporary_target: bool = False,
     ):
         if target == ".":
             target = os.path.basename(source)
             if not target:
                 target = os.path.basename(os.path.dirname(source))
 
-        temporary_target = {
+        temporary_target: Dict[Any, Any] = {
             False: None,
             True: "{dirname}-downloading/{basename}",
             "-tmp/*": "{dirname}-tmp/{basename}",
@@ -223,8 +512,14 @@ class Transfer:
         cls = _find_transfer_class(self.source, self.target)
         self.loader = cls()
 
-    def run(self):
+    def run(self) -> "Transfer":
+        """Execute the transfer process.
 
+        Returns
+        -------
+        Transfer
+            The Transfer instance.
+        """
         target = self.loader.get_temporary_target(self.target, self.temporary_target)
         if target != self.target:
             LOGGER.info(f"Using temporary target {target} to copy to {self.target}")
@@ -256,16 +551,51 @@ class Transfer:
 
         return self
 
-    def rename_target(self, target, new_target):
+    def rename_target(self, target: str, new_target: str) -> None:
+        """Rename the target to a new target path.
+
+        Parameters
+        ----------
+        target : str
+            The original target path.
+        new_target : str
+            The new target path.
+        """
         if target != new_target:
             LOGGER.info(f"Renaming temporary target {target} into {self.target}")
             return self.loader.rename_target(target, new_target)
 
-    def delete_target(self, target):
+    def delete_target(self, target: str) -> None:
+        """Delete the target if it exists.
+
+        Parameters
+        ----------
+        target : str
+            The target path.
+        """
         return self.loader.delete_target(target)
 
 
-def _find_transfer_class(source, target):
+def _find_transfer_class(source: str, target: str) -> type:
+    """Find the appropriate transfer class based on the source and target locations.
+
+    Parameters
+    ----------
+    source : str
+        The source location.
+    target : str
+        The target location.
+
+    Returns
+    -------
+    type
+        The transfer class.
+
+    Raises
+    ------
+    TransferMethodNotImplementedError
+        If the transfer method is not implemented.
+    """
     from_ssh = source.startswith("ssh://")
     into_ssh = target.startswith("ssh://")
 
@@ -298,8 +628,12 @@ def _find_transfer_class(source, target):
 
 
 # this is the public API
-def transfer(*args, **kwargs) -> Loader:
-    """Parameters
+def transfer(
+    source, target, *, overwrite=False, resume=False, verbosity=1, progress=None, threads=1, temporary_target=False
+) -> Loader:
+    """Transfer files or folders from the source to the target location.
+
+    Parameters
     ----------
     source : str
         A path to a local file or folder or a URL to a file or a folder on S3.
@@ -316,7 +650,7 @@ def transfer(*args, **kwargs) -> Loader:
         By default False
     verbosity : int, optional
         The level of verbosity, by default 1
-    progress: callable, optional
+    progress : callable, optional
         A callable that will be called with the number of files, the total size of the files, the total size
         transferred and a boolean indicating if the transfer has started. By default None
     threads : int, optional
@@ -325,8 +659,22 @@ def transfer(*args, **kwargs) -> Loader:
         Experimental feature
         If True and if the target location supports it, the data will be uploaded to a temporary location
         then renamed to the final location. Supported by SSH and local targets, not supported by S3.
-        By default False
+        By default False.
+
+    Returns
+    -------
+    Loader
+        The Loader instance.
     """
-    copier = Transfer(*args, **kwargs)
+    copier = Transfer(
+        source=source,
+        target=target,
+        overwrite=overwrite,
+        resume=resume,
+        verbosity=verbosity,
+        progress=progress,
+        threads=threads,
+        temporary_target=temporary_target,
+    )
     copier.run()
     return copier