anemoi-utils 0.3.4__tar.gz → 0.3.6__tar.gz

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/.pre-commit-config.yaml

@@ -38,7 +38,7 @@ repos:
   hooks:
   - id: ruff
     # Next line if for documenation cod snippets
-    exclude: '^(dev/.*|[A-Za-z].*_)\.py$'
+    exclude: '^[^_].*_\.py$'
     args:
     - --line-length=120
     - --fix

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.3.4
+Version: 0.3.6
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/docs/conf.py

@@ -93,12 +93,16 @@ intersphinx_mapping = {
         "https://anemoi-inference.readthedocs.io/en/latest/",
         ("../../anemoi-inference/docs/_build/html/objects.inv", None),
     ),
+    "anemoi-graphs": (
+        "https://anemoi-graphs.readthedocs.io/en/latest/",
+        ("../../anemoi-graphs/docs/_build/html/objects.inv", None),
+    ),
+    "anemoi-registry": (
+        "https://anemoi-registry.readthedocs.io/en/latest/",
+        ("../../anemoi-registry/docs/_build/html/objects.inv", None),
+    ),
 }
 
-
-# https://www.notion.so/Deepnote-Launch-Buttons-63c642a5e875463495ed2341e83a4b2a
-
-
 # -- Options for HTML output -------------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/docs/index.rst

@@ -47,8 +47,10 @@ of the *Anemoi* packages.
 -  :ref:`anemoi-utils <anemoi-utils:index-page>`
 -  :ref:`anemoi-datasets <anemoi-datasets:index-page>`
 -  :ref:`anemoi-models <anemoi-models:index-page>`
+-  :ref:`anemoi-graphs <anemoi-graphs:index-page>`
 -  :ref:`anemoi-training <anemoi-training:index-page>`
 -  :ref:`anemoi-inference <anemoi-inference:index-page>`
+-  :ref:`anemoi-registry <anemoi-registry:index-page>`
 
 *********
  License

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/src/anemoi/utils/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.3.4'
-__version_tuple__ = version_tuple = (0, 3, 4)
+__version__ = version = '0.3.6'
+__version_tuple__ = version_tuple = (0, 3, 6)

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/src/anemoi/utils/cli.py

@@ -64,7 +64,7 @@ def register_commands(here, package, select, fail=None):
     result = {}
     not_available = {}
 
-    for p in os.listdir(here):
+    for p in sorted(os.listdir(here)):
         full = os.path.join(here, p)
         if p.startswith("_"):
             continue

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/src/anemoi/utils/config.py

@@ -6,8 +6,12 @@
 # nor does it submit to any jurisdiction.
 
 
+import json
 import logging
 import os
+import threading
+
+import yaml
 
 try:
     import tomllib  # Only available since 3.11
@@ -40,10 +44,49 @@ class DotDict(dict):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+
         for k, v in self.items():
             if isinstance(v, dict):
                 self[k] = DotDict(v)
 
+            if isinstance(v, list):
+                self[k] = [DotDict(i) if isinstance(i, dict) else i for i in v]
+
+            if isinstance(v, tuple):
+                self[k] = [DotDict(i) if isinstance(i, dict) else i for i in v]
+
+    @classmethod
+    def from_file(cls, path: str):
+        _, ext = os.path.splitext(path)
+        if ext == ".yaml" or ext == ".yml":
+            return cls.from_yaml_file(path)
+        elif ext == ".json":
+            return cls.from_json_file(path)
+        elif ext == ".toml":
+            return cls.from_toml_file(path)
+        else:
+            raise ValueError(f"Unknown file extension {ext}")
+
+    @classmethod
+    def from_yaml_file(cls, path: str):
+        with open(path, "r") as file:
+            data = yaml.safe_load(file)
+
+        return cls(data)
+
+    @classmethod
+    def from_json_file(cls, path: str):
+        with open(path, "r") as file:
+            data = json.load(file)
+
+        return cls(data)
+
+    @classmethod
+    def from_toml_file(cls, path: str):
+        with open(path, "r") as file:
+            data = tomllib.load(file)
+        return cls(data)
+
     def __getattr__(self, attr):
         try:
             return self[attr]
@@ -60,16 +103,10 @@ class DotDict(dict):
 
 
 CONFIG = None
+CONFIG_LOCK = threading.Lock()
 
 
-def load_config():
-    """Load the configuration from `~/.anemoi.toml`.
-
-    Returns
-    -------
-    DotDict
-        The configuration
-    """
+def _load_config():
     global CONFIG
     if CONFIG is not None:
         return CONFIG
@@ -86,9 +123,20 @@ def load_config():
     return DotDict(CONFIG)
 
 
-def save_config():
-    """Save the configuration to `~/.anemoi.toml`."""
+def load_config():
+    """Load the configuration from `~/.anemoi.toml`.
+
+    Returns
+    -------
+    DotDict
+        The configuration
+    """
+    with CONFIG_LOCK:
+        return _load_config()
+
 
+def check_config_mode():
     conf = os.path.expanduser("~/.anemoi.toml")
-    with open(conf, "w") as f:
-        tomllib.dump(CONFIG, f)
+    mode = os.stat(conf).st_mode
+    if mode & 0o777 != 0o600:
+        raise SystemError(f"Configuration file {conf} is not secure. " "Please run `chmod 600 ~/.anemoi.toml`.")

anemoi_utils-0.3.6/src/anemoi/utils/s3.py

@@ -0,0 +1,463 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""This module provides functions to upload, download, list and delete files and folders on S3.
+The functions of this package expect that the AWS credentials are set up in the environment
+typicaly by setting the `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables or
+by creating a `~/.aws/credentials` file. It is also possible to set the `endpoint_url` in the same file
+to use a different S3 compatible service::
+
+    [default]
+    endpoint_url = https://some-storage.somewhere.world
+    aws_access_key_id = xxxxxxxxxxxxxxxxxxxxxxxx
+    aws_secret_access_key = xxxxxxxxxxxxxxxxxxxxxxxx
+
+"""
+
+import concurrent
+import logging
+import os
+import threading
+from copy import deepcopy
+
+import tqdm
+
+from .config import check_config_mode
+from .config import load_config
+from .humanize import bytes
+
+LOGGER = logging.getLogger(__name__)
+
+
+# s3_clients are not thread-safe, so we need to create a new client for each thread
+
+thread_local = threading.local()
+
+
+def s3_client(bucket):
+    import boto3
+
+    config = load_config()
+    if "object-storage" in config:
+        check_config_mode()
+
+    if not hasattr(thread_local, "s3_clients"):
+        thread_local.s3_clients = {}
+
+    if bucket not in thread_local.s3_clients:
+
+        options = {}
+        options.update(config.get("object-storage", {}))
+        options.update(config.get("object-storage", {}).get(bucket, {}))
+
+        type = options.pop("type", "s3")
+        if type != "s3":
+            raise ValueError(f"Unsupported object storage type {type}")
+
+        if "config" in options:
+            from botocore.client import Config
+
+            options["config"] = Config(**options["config"])
+            del options["config"]
+
+        thread_local.s3_clients[bucket] = boto3.client("s3", **options)
+
+    return thread_local.s3_clients[bucket]
+
+
+class Transfer:
+
+    def transfer_folder(self, *, source, target, overwrite=False, resume=False, verbosity=1, threads=1):
+        assert verbosity == 1, verbosity
+
+        # from boto3.s3.transfer import TransferConfig
+        # config = TransferConfig(use_threads=False)
+        config = None
+        with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
+            try:
+                if verbosity > 0:
+                    LOGGER.info(f"{self.action} {source} to {target}")
+
+                total = 0
+
+                futures = []
+                for name in self.list_source(source):
+
+                    futures.append(
+                        executor.submit(
+                            self.transfer_file,
+                            source=self.source_path(name, source),
+                            target=self.target_path(name, source, target),
+                            overwrite=overwrite,
+                            resume=resume,
+                            verbosity=verbosity - 1,
+                            config=config,
+                        )
+                    )
+                    total += self.source_size(name)
+
+                    if len(futures) % 10000 == 0:
+                        if verbosity > 0:
+                            LOGGER.info(f"Preparing transfer, {len(futures):,} files... ({bytes(total)})")
+                        done, _ = concurrent.futures.wait(
+                            futures,
+                            timeout=0.001,
+                            return_when=concurrent.futures.FIRST_EXCEPTION,
+                        )
+                        # Trigger exceptions if any
+                        for future in done:
+                            future.result()
+
+                if verbosity > 0:
+                    LOGGER.info(f"{self.action} {len(futures):,} files ({bytes(total)})")
+                    with tqdm.tqdm(total=total, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+                        for future in futures:
+                            pbar.update(future.result())
+                else:
+                    for future in futures:
+                        future.result()
+
+            except Exception:
+                executor.shutdown(wait=False, cancel_futures=True)
+                raise
+
+
+class Upload(Transfer):
+    action = "Uploading"
+
+    def list_source(self, source):
+        for root, _, files in os.walk(source):
+            for file in files:
+                yield os.path.join(root, file)
+
+    def source_path(self, local_path, source):
+        return local_path
+
+    def target_path(self, source_path, source, target):
+        relative_path = os.path.relpath(source_path, source)
+        s3_path = os.path.join(target, relative_path)
+        return s3_path
+
+    def source_size(self, local_path):
+        return os.path.getsize(local_path)
+
+    def transfer_file(self, source, target, overwrite, resume, verbosity, config=None):
+
+        from botocore.exceptions import ClientError
+
+        assert target.startswith("s3://")
+
+        _, _, bucket, key = target.split("/", 3)
+        s3 = s3_client(bucket)
+
+        size = os.path.getsize(source)
+
+        if verbosity > 0:
+            LOGGER.info(f"{self.action} {source} to {target} ({bytes(size)})")
+
+        try:
+            results = s3.head_object(Bucket=bucket, Key=key)
+            remote_size = int(results["ContentLength"])
+        except ClientError as e:
+            if e.response["Error"]["Code"] != "404":
+                raise
+            remote_size = None
+
+        if remote_size is not None:
+            if remote_size != size:
+                LOGGER.warning(
+                    f"{target} already exists, but with different size, re-uploading (remote={remote_size}, local={size})"
+                )
+            elif resume:
+                # LOGGER.info(f"{target} already exists, skipping")
+                return size
+
+        if remote_size is not None and not overwrite and not resume:
+            raise ValueError(f"{target} already exists, use 'overwrite' to replace or 'resume' to skip")
+
+        if verbosity > 0:
+            with tqdm.tqdm(total=size, unit="B", unit_scale=True, unit_divisor=1024, leave=False) as pbar:
+                s3.upload_file(source, bucket, key, Callback=lambda x: pbar.update(x), Config=config)
+        else:
+            s3.upload_file(source, bucket, key, Config=config)
+
+        return size
+
+
+class Download(Transfer):
+    action = "Downloading"
+
+    def list_source(self, source):
+        yield from _list_objects(source)
+
+    def source_path(self, s3_object, source):
+        _, _, bucket, _ = source.split("/", 3)
+        return f"s3://{bucket}/{s3_object['Key']}"
+
+    def target_path(self, s3_object, source, target):
+        _, _, _, folder = source.split("/", 3)
+        local_path = os.path.join(target, os.path.relpath(s3_object["Key"], folder))
+        os.makedirs(os.path.dirname(local_path), exist_ok=True)
+        return local_path
+
+    def source_size(self, s3_object):
+        return s3_object["Size"]
+
+    def transfer_file(self, source, target, overwrite, resume, verbosity, config=None):
+        # from boto3.s3.transfer import TransferConfig
+
+        _, _, bucket, key = source.split("/", 3)
+        s3 = s3_client(bucket)
+
+        try:
+            response = s3.head_object(Bucket=bucket, Key=key)
+        except s3.exceptions.ClientError as e:
+            print(e.response["Error"]["Code"], e.response["Error"]["Message"], bucket, key)
+            if e.response["Error"]["Code"] == "404":
+                raise ValueError(f"{source} does not exist ({bucket}, {key})")
+            raise
+
+        size = int(response["ContentLength"])
+
+        if verbosity > 0:
+            LOGGER.info(f"Downloading {source} to {target} ({bytes(size)})")
+
+        if overwrite:
+            resume = False
+
+        if resume:
+            if os.path.exists(target):
+                local_size = os.path.getsize(target)
+                if local_size != size:
+                    LOGGER.warning(
+                        f"{target} already with different size, re-downloading (remote={size}, local={size})"
+                    )
+                else:
+                    # if verbosity > 0:
+                    #     LOGGER.info(f"{target} already exists, skipping")
+                    return size
+
+        if os.path.exists(target) and not overwrite:
+            raise ValueError(f"{target} already exists, use 'overwrite' to replace or 'resume' to skip")
+
+        if verbosity > 0:
+            with tqdm.tqdm(total=size, unit="B", unit_scale=True, unit_divisor=1024, leave=False) as pbar:
+                s3.download_file(bucket, key, target, Callback=lambda x: pbar.update(x), Config=config)
+        else:
+            s3.download_file(bucket, key, target, Config=config)
+
+        return size
+
+
+def upload(source, target, *, overwrite=False, resume=False, verbosity=1, threads=1):
+    """Upload a file or a folder to S3.
+
+    Parameters
+    ----------
+    source : str
+        A path to a file or a folder to upload.
+    target : str
+        A URL to a file or a folder on S3. The url should start with 's3://'.
+    overwrite : bool, optional
+        If the data is alreay on S3 it will be overwritten, by default False
+    resume : bool, optional
+        If the data is alreay on S3 it will not be uploaded, unless the remote file
+        has a different size, by default False
+    threads : int, optional
+        The number of threads to use when uploading a directory, by default 1
+    """
+
+    uploader = Upload()
+    if os.path.isdir(source):
+        uploader.transfer_folder(
+            source=source, target=target, overwrite=overwrite, resume=resume, verbosity=verbosity, threads=threads
+        )
+    else:
+        uploader.transfer_file(source=source, target=target, overwrite=overwrite, resume=resume, verbosity=verbosity)
+
+
+def download(source, target, *, overwrite=False, resume=False, verbosity=1, threads=1):
+    """Download a file or a folder from S3.
+
+    Parameters
+    ----------
+    source : str
+        The URL of a file or a folder on S3. The url should start with 's3://'. If the URL ends with a '/' it is
+        assumed to be a folder, otherwise it is assumed to be a file.
+    target : str
+        The local path where the file or folder will be downloaded.
+    overwrite : bool, optional
+        If false, files which have already been download will be skipped, unless their size
+        does not match their size on S3 , by default False
+    resume : bool, optional
+        If the data is alreay on local it will not be downloaded, unless the remote file
+        has a different size, by default False
+    threads : int, optional
+        The number of threads to use when downloading a directory, by default 1
+    """
+    assert source.startswith("s3://")
+
+    downloader = Download()
+
+    if source.endswith("/"):
+        downloader.transfer_folder(
+            source=source,
+            target=target,
+            overwrite=overwrite,
+            resume=resume,
+            verbosity=verbosity,
+            threads=threads,
+        )
+    else:
+        downloader.transfer_file(
+            source=source,
+            target=target,
+            overwrite=overwrite,
+            resume=resume,
+            verbosity=verbosity,
+        )
+
+
+def _list_objects(target, batch=False):
+    _, _, bucket, prefix = target.split("/", 3)
+    s3 = s3_client(bucket)
+
+    paginator = s3.get_paginator("list_objects_v2")
+
+    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
+        if "Contents" in page:
+            objects = deepcopy(page["Contents"])
+            if batch:
+                yield objects
+            else:
+                yield from objects
+
+
+def _delete_folder(target):
+    _, _, bucket, _ = target.split("/", 3)
+    s3 = s3_client(bucket)
+
+    total = 0
+    for batch in _list_objects(target, batch=True):
+        LOGGER.info(f"Deleting {len(batch):,} objects from {target}")
+        s3.delete_objects(Bucket=bucket, Delete={"Objects": [{"Key": o["Key"]} for o in batch]})
+        total += len(batch)
+        LOGGER.info(f"Deleted {len(batch):,} objects (total={total:,})")
+
+
+def _delete_file(target):
+    from botocore.exceptions import ClientError
+
+    _, _, bucket, key = target.split("/", 3)
+    s3 = s3_client(bucket)
+
+    try:
+        s3.head_object(Bucket=bucket, Key=key)
+        exits = True
+    except ClientError as e:
+        if e.response["Error"]["Code"] != "404":
+            raise
+        exits = False
+
+    if not exits:
+        LOGGER.warning(f"{target} does not exist. Did you mean to delete a folder? Then add a trailing '/'")
+        return
+
+    LOGGER.info(f"Deleting {target}")
+    print(s3.delete_object(Bucket=bucket, Key=key))
+    LOGGER.info(f"{target} is deleted")
+
+
+def delete(target):
+    """Delete a file or a folder from S3.
+
+    Parameters
+    ----------
+    target : str
+        The URL of a file or a folder on S3. The url should start with 's3://'. If the URL ends with a '/' it is
+        assumed to be a folder, otherwise it is assumed to be a file.
+    """
+
+    assert target.startswith("s3://")
+
+    if target.endswith("/"):
+        _delete_folder(target)
+    else:
+        _delete_file(target)
+
+
+def list_folder(folder):
+    """List the sub folders in a folder on S3.
+
+    Parameters
+    ----------
+    folder : str
+        The URL of a folder on S3. The url should start with 's3://'.
+
+    Returns
+    -------
+    list
+        A list of the subfolders names in the folder.
+    """
+
+    assert folder.startswith("s3://")
+    if not folder.endswith("/"):
+        folder += "/"
+
+    _, _, bucket, prefix = folder.split("/", 3)
+
+    s3 = s3_client(bucket)
+    paginator = s3.get_paginator("list_objects_v2")
+
+    for page in paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/"):
+        if "CommonPrefixes" in page:
+            yield from [folder + _["Prefix"] for _ in page.get("CommonPrefixes")]
+
+
+def object_info(target):
+    """Get information about an object on S3.
+
+    Parameters
+    ----------
+    target : str
+        The URL of a file or a folder on S3. The url should start with 's3://'.
+
+    Returns
+    -------
+    dict
+        A dictionary with information about the object.
+    """
+
+    _, _, bucket, key = target.split("/", 3)
+    s3 = s3_client(bucket)
+
+    try:
+        return s3.head_object(Bucket=bucket, Key=key)
+    except s3.exceptions.ClientError as e:
+        if e.response["Error"]["Code"] == "404":
+            raise ValueError(f"{target} does not exist")
+        raise
+
+
+def object_acl(target):
+    """Get information about an object's ACL on S3.
+
+    Parameters
+    ----------
+    target : str
+        The URL of a file or a folder on S3. The url should start with 's3://'.
+
+    Returns
+    -------
+    dict
+        A dictionary with information about the object's ACL.
+    """
+
+    _, _, bucket, key = target.split("/", 3)
+    s3 = s3_client()
+
+    return s3.get_object_acl(Bucket=bucket, Key=key)

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/src/anemoi_utils.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.3.4
+Version: 0.3.6
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License

{anemoi_utils-0.3.4 → anemoi_utils-0.3.6}/tests/test_utils.py

@@ -12,7 +12,7 @@ from anemoi.utils.grib import shortname_to_paramid
 
 
 def test_dotdict():
-    d = DotDict(a=1, b=2, c=dict(d=3, e=4))
+    d = DotDict(a=1, b=2, c=dict(d=3, e=4), e=[1, dict(a=3), 3])
     assert d.a == 1
     assert d.b == 2
     assert d.c.d == 3
@@ -27,6 +27,8 @@ def test_dotdict():
     d.d.x = 6
     assert d.d.x == 6
 
+    assert d.e[1].a == 3
+
 
 def test_grib():
     assert shortname_to_paramid("2t") == 167

anemoi_utils-0.3.4/src/anemoi/utils/s3.py

@@ -1,280 +0,0 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
-# This software is licensed under the terms of the Apache Licence Version 2.0
-# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
-# In applying this licence, ECMWF does not waive the privileges and immunities
-# granted to it by virtue of its status as an intergovernmental organisation
-# nor does it submit to any jurisdiction.
-
-"""This module provides functions to upload, download, list and delete files and folders on S3.
-The functions of this package expect that the AWS credentials are set up in the environment
-typicaly by setting the `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables or
-by creating a `~/.aws/credentials` file. It is also possible to set the `endpoint_url` in the same file
-to use a different S3 compatible service::
-
-    [default]
-    endpoint_url = https://some-storage.somewhere.world
-    aws_access_key_id = xxxxxxxxxxxxxxxxxxxxxxxx
-    aws_secret_access_key = xxxxxxxxxxxxxxxxxxxxxxxx
-
-"""
-
-import logging
-import os
-import threading
-from concurrent.futures import ThreadPoolExecutor
-from contextlib import closing
-
-import tqdm
-
-LOGGER = logging.getLogger(__name__)
-
-
-# s3_clients are not thread-safe, so we need to create a new client for each thread
-
-thread_local = threading.local()
-
-
-def _s3_client():
-    import boto3
-
-    if not hasattr(thread_local, "s3_client"):
-        thread_local.s3_client = boto3.client("s3")
-    return thread_local.s3_client
-
-
-def _upload_file(source, target, overwrite=False, ignore_existing=False):
-    from botocore.exceptions import ClientError
-
-    assert target.startswith("s3://")
-
-    _, _, bucket, key = target.split("/", 3)
-
-    LOGGER.info(f"Uploading {source} to {target}")
-    s3_client = _s3_client()
-
-    size = os.path.getsize(source)
-    try:
-        results = s3_client.head_object(Bucket=bucket, Key=key)
-        remote_size = int(results["ContentLength"])
-    except ClientError as e:
-        if e.response["Error"]["Code"] != "404":
-            raise
-        remote_size = None
-
-    if remote_size is not None:
-        if remote_size != size:
-            LOGGER.warning(f"{target} already exists, but with different size, re-uploading")
-            overwrite = True
-
-        if ignore_existing:
-            LOGGER.info(f"{target} already exists, skipping")
-            return
-
-    if remote_size is not None and not overwrite:
-        raise ValueError(f"{target} already exists, use 'overwrite' to replace or 'ignore_existing' to skip")
-
-    with closing(tqdm.tqdm(total=size, unit="B", unit_scale=True, leave=False)) as t:
-        s3_client.upload_file(source, bucket, key, Callback=lambda x: t.update(x))
-
-
-def _local_file_list(source):
-    for root, _, files in os.walk(source):
-        for file in files:
-            yield os.path.join(root, file)
-
-
-def _upload_folder(source, target, overwrite=False, ignore_existing=False, threads=1):
-    total = sum(1 for _ in _local_file_list(source))
-
-    with ThreadPoolExecutor(max_workers=threads) as executor:
-        futures = []
-        for local_path in _local_file_list(source):
-            relative_path = os.path.relpath(local_path, source)
-            s3_path = os.path.join(target, relative_path)
-            futures.append(executor.submit(_upload_file, local_path, s3_path, overwrite, ignore_existing))
-
-        for future in tqdm.tqdm(futures, total=total):
-            future.result()
-
-
-def upload(source, target, overwrite=False, ignore_existing=False, threads=1, show_progress=True):
-    """Upload a file or a folder to S3.
-
-    Parameters
-    ----------
-    source : str
-        A path to a file or a folder to upload.
-    target : str
-        A URL to a file or a folder on S3. The url should start with 's3://'.
-    overwrite : bool, optional
-        If the data is alreay on S3 it will be overwritten, by default False
-    ignore_existing : bool, optional
-        If the data is alreay on S3 it will not be uploaded, unless the remote file
-        has a different size, by default False
-    threads : int, optional
-        The number of threads to use when uploading a directory, by default 1
-    """
-    if os.path.isdir(source):
-        _upload_folder(source, target, overwrite, ignore_existing, threads)
-    else:
-        _upload_file(source, target, overwrite, ignore_existing)
-
-
-def _download_file(source, target, overwrite=False, ignore_existing=False):
-    s3_client = _s3_client()
-    _, _, bucket, key = source.split("/", 3)
-
-    response = s3_client.head_object(Bucket=bucket, Key=key)
-    size = int(response["ContentLength"])
-
-    if os.path.exists(target):
-
-        if os.path.exists(target) and os.path.getsize(target) != size:
-            LOGGER.info(f"{target} already with different size, re-downloading")
-            overwrite = True
-
-        if not overwrite and not ignore_existing:
-            raise ValueError(f"{target} already exists, use 'overwrite' to replace or 'ignore_existing' to skip")
-
-        if ignore_existing:
-            LOGGER.info(f"{target} already exists, skipping")
-            return
-
-    with closing(tqdm.tqdm(total=size, unit="B", unit_scale=True, leave=False)) as t:
-        s3_client.download_file(bucket, key, target, Callback=lambda x: t.update(x))
-
-
-def _download_folder(source, target, overwrite=False, ignore_existing=False, threads=1):
-    source = source.rstrip("/")
-    _, _, bucket, folder = source.split("/", 3)
-    total = _count_objects_in_folder(source)
-
-    with ThreadPoolExecutor(max_workers=threads) as executor:
-        futures = []
-        for o in _list_folder(source):
-            name = o["Key"]
-            local_path = os.path.join(target, os.path.relpath(name, folder))
-            os.makedirs(os.path.dirname(local_path), exist_ok=True)
-            futures.append(
-                executor.submit(_download_file, f"s3://{bucket}/{name}", local_path, overwrite, ignore_existing)
-            )
-
-        for future in tqdm.tqdm(futures, total=total):
-            future.result()
-
-
-def download(source, target, overwrite=False, ignore_existing=False, threads=1, show_progress=True):
-    """Download a file or a folder from S3.
-
-    Parameters
-    ----------
-    source : str
-        The URL of a file or a folder on S3. The url should start with 's3://'. If the URL ends with a '/' it is
-        assumed to be a folder, otherwise it is assumed to be a file.
-    target : str
-        The local path where the file or folder will be downloaded.
-    overwrite : bool, optional
-        If false, files which have already been download will be skipped, unless their size
-        does not match their size on S3 , by default False
-    ignore_existing : bool, optional
-        If the data is alreay on local it will not be downloaded, unless the remote file
-        has a different size, by default False
-    threads : int, optional
-        The number of threads to use when downloading a directory, by default 1
-    """
-    assert source.startswith("s3://")
-
-    if source.endswith("/"):
-        _download_folder(source, target, overwrite, ignore_existing, threads)
-    else:
-        _download_file(source, target, overwrite, ignore_existing)
-
-
-def _list_folder(target, batch=False):
-    s3_client = _s3_client()
-    _, _, bucket, prefix = target.split("/", 3)
-
-    paginator = s3_client.get_paginator("list_objects_v2")
-    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
-        if "Contents" in page:
-            objects = [{"Key": obj["Key"]} for obj in page["Contents"]]
-            if batch:
-                yield objects
-            else:
-                yield from objects
-
-
-def _count_objects_in_folder(target):
-    return sum(len(_) for _ in _list_folder(target, batch=True))
-
-
-def _delete_folder(target, threads):
-    s3_client = _s3_client()
-    _, _, bucket, _ = target.split("/", 3)
-
-    for batch in _list_folder(target, batch=True):
-        s3_client.delete_objects(Bucket=bucket, Delete={"Objects": batch})
-        LOGGER.info(f"Deleted {len(batch)} objects")
-
-
-def _delete_file(target):
-    s3_client = _s3_client()
-    _, _, bucket, key = target.split("/", 3)
-
-    LOGGER.info(f"Deleting {target}")
-    s3_client.delete_object(Bucket=bucket, Key=key)
-    LOGGER.info(f"{target} is deleted")
-
-
-def delete(target):
-    """Delete a file or a folder from S3.
-
-    Parameters
-    ----------
-    target : str
-        The URL of a file or a folder on S3. The url should start with 's3://'. If the URL ends with a '/' it is
-        assumed to be a folder, otherwise it is assumed to be a file.
-    """
-
-    assert target.startswith("s3://")
-
-    if target.endswith("/"):
-        _delete_folder(target)
-    else:
-        _delete_file(target)
-
-
-def list_folder(target):
-    """List the objects in a folder on S3.
-
-    Parameters
-    ----------
-    target : str
-        The URL of a folder on S3. The url should start with 's3://'.
-
-    Returns
-    -------
-    list
-        A list of the objects names in the folder.
-    """
-
-    assert target.startswith("s3://")
-    return [o["Key"] for o in _list_folder(target)]
-
-
-def count_objects_in_folder(target):
-    """Count the objects in a folder on S3.
-
-    Parameters
-    ----------
-    target : str
-        The URL of a folder on S3. The url should start with 's3://'.
-
-    Returns
-    -------
-    int
-        The number of objects in the folder.
-    """
-
-    assert target.startswith("s3://")
-    return _count_objects_in_folder(target)