Refactor LargeFile

great_ai/src/great_ai/large_file/__init__.py

@@ -0,0 +1 @@
+from .large_file import LargeFile, LargeFileS3

great_ai/src/great_ai/large_file/__main__.py

@@ -4,16 +4,30 @@ from pathlib import Path
 
 from great_ai.utilities.logger import get_logger
 
-from .large_file import LargeFile
+from .large_file import LargeFileLocal, LargeFileS3
 from .parse_arguments import parse_arguments
 
-logger = get_logger("open_s3")
+logger = get_logger("large_file")
 
 
 def main() -> None:
     parser, args = parse_arguments()
 
-    LargeFile.configure_credentials_from_file(args.secrets)
+    factory = {"s3": LargeFileS3, "local": LargeFileLocal}
+
+    if args.driver not in factory:
+        raise ValueError(
+            f"Driver {args.driver} does not exits, available options: {' ,'.join(factory.keys())}"
+        )
+
+    large_file = factory[args.driver]
+
+    if args.driver != "local":
+        if args.secrets is None:
+            raise ValueError(
+                "Providing a credentials file is required when the driver mode is not `local`."
+            )
+        large_file.configure_credentials_from_file(args.secrets)  # type: ignore
 
     if not args.cache and not args.push and not args.delete:
         logger.warning("No action required.")
@@ -24,16 +38,16 @@ def main() -> None:
             split = c.split(":")
             file_name = split[0]
             version = None if len(split) == 1 else int(split[1])
-            LargeFile(file_name, "r", version=version).get()
+            large_file(file_name, "r", version=version).get()
 
     if args.push:
         for p in args.push:
             path = Path(p)
-            LargeFile(path.name, "w").push(path)
+            large_file(path.name, "w").push(path)
 
     if args.delete:
         for f in args.delete:
-            LargeFile(f).delete()
+            large_file(f).delete()
 
 
 if __name__ == "__main__":

great_ai/src/great_ai/large_file/large_file/__init__.py

@@ -0,0 +1,3 @@
+from .large_file import LargeFile
+from .large_file_local import LargeFileLocal
+from .large_file_s3 import LargeFileS3

great_ai/src/great_ai/large_file/large_file/large_file.py

@@ -0,0 +1,300 @@
+import os
+import shutil
+import tempfile
+from abc import ABC, abstractmethod
+from pathlib import Path
+from types import TracebackType
+from typing import IO, Any, List, Optional, Type, Union, cast
+
+from great_ai.utilities.logger import get_logger
+
+from ..helper import human_readable_to_byte
+from ..models import DataInstance
+
+logger = get_logger("large_file")
+
+
+CACHE_NAME_VERSION_SEPARATOR = "-"
+COMPRESSION_ALGORITHM = "gztar"
+ARCHIVE_EXTENSION = ".tar.gz"
+
+
+class LargeFile(ABC):
+    """
+    Store large files remotely. Use local cache for speed up.
+
+    Examples:
+
+    ```
+    with LargeFile("test.txt", "w", keep_last_n=3) as f:
+        for i in range(1000000):
+            f.write('test\n')
+
+    with LargeFile("test.txt", "r") as f:
+        print(f.readlines()[0])
+
+    path_to_cached_text_file = LargeFile("test.txt", version=0).get()
+    ```
+
+    By default, files are stored in the ".cache" folder and the
+    least recently use is deleted after the overall size reaches 30 GBs.
+
+    Change it with the following properties.
+
+    ```
+    LargeFile.cache_path = Path(".cache")
+    LargeFile.max_cache_size = "30GB"
+    ```
+    """
+
+    cache_path = Path(".cache")
+    max_cache_size: Optional[str] = "30GB"
+
+    def __init__(
+        self,
+        name: str,
+        mode: str = "r",
+        *,
+        buffering: int = -1,
+        encoding: Optional[str] = None,
+        errors: Optional[str] = None,
+        newline: Optional[str] = None,
+        version: Optional[int] = None,
+        keep_last_n: Optional[int] = None,
+        cache_only_mode: bool = False,
+    ):
+        self._name = name
+        self._version = version
+        self._mode = mode
+        self._keep_last_n = keep_last_n
+        self._cache_only_mode = cache_only_mode
+
+        self._buffering = buffering
+        self._encoding = encoding
+        self._errors = errors
+        self._newline = newline
+
+        LargeFile.cache_path.mkdir(parents=True, exist_ok=True)
+
+        self._find_instances()
+        self._check_mode_and_set_version()
+
+    def __enter__(self) -> IO:
+        self._file: IO[Any] = (
+            tempfile.NamedTemporaryFile(
+                mode=self._mode,
+                buffering=self._buffering,
+                encoding=self._encoding,
+                newline=self._newline,
+                errors=self._errors,
+                delete=False,
+                prefix="large_file-",
+            )
+            if "w" in self._mode
+            else open(
+                self.get(),
+                mode=self._mode,
+                buffering=self._buffering,
+                encoding=self._encoding,
+                newline=self._newline,
+                errors=self._errors,
+            )
+        )
+
+        return self._file
+
+    def __exit__(
+        self,
+        type: Optional[Type[BaseException]],
+        exc: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        self._file.close()
+
+        if type is None:
+            if "w" in self._mode:
+                self.push(Path(self._file.name))
+                os.unlink(self._file.name)
+        else:
+            logger.exception("Could not finish operation.")
+
+        return True
+
+    @property
+    def _local_name(self) -> str:
+        return f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}{self.version}"
+
+    @property
+    def version(self) -> int:
+        return cast(int, self._version)
+
+    def get(self, hide_progress: bool = False) -> Path:
+        remote_path = next(
+            i.remote_path for i in self._instances if i.version == self._version
+        )
+
+        destination = self.cache_path / self._local_name
+        if not destination.exists():
+            logger.info(f"File {self._local_name} does not exist locally")
+
+            with tempfile.TemporaryDirectory() as tmp:
+                local_root_path = Path(tmp)
+                tmp_file_archive = (
+                    local_root_path / f"{self._local_name}{ARCHIVE_EXTENSION}"
+                )
+                self._download(
+                    remote_path, tmp_file_archive, hide_progress=hide_progress
+                )
+
+                logger.info(f"Decompressing {self._local_name}")
+                shutil.unpack_archive(str(tmp_file_archive), tmp, COMPRESSION_ALGORITHM)
+                shutil.move(str(local_root_path / self._local_name), str(destination))
+        else:
+            logger.info(f"File {self._local_name} found in cache")
+
+        return destination
+
+    def push(self, path: Union[Path, str], hide_progress: bool = False) -> None:
+        if isinstance(path, str):
+            path = Path(path)
+
+        with tempfile.TemporaryDirectory() as tmp:
+            if path.is_file():
+                logger.info(f"Copying file for {self._local_name}")
+                copy: Any = shutil.copy
+            else:
+                logger.info(f"Copying directory for {self._local_name}")
+                copy = shutil.copytree
+
+            try:
+                # Make local copy in the cache
+                copy(str(path), str(self.cache_path / self._local_name))
+            except shutil.SameFileError:
+                pass  # No worries
+
+            copy(str(path), str(Path(tmp) / self._local_name))
+
+            with tempfile.TemporaryDirectory() as tmp2:
+                # A directory has to be zipped and it cannot contain the output of the zipping
+                logger.info(f"Compressing {self._local_name}")
+                shutil.make_archive(
+                    str(Path(tmp2) / self._local_name),
+                    COMPRESSION_ALGORITHM,
+                    tmp,
+                )
+
+                file_to_be_uploaded = (
+                    Path(tmp2) / f"{self._local_name}{ARCHIVE_EXTENSION}"
+                )
+                self._upload(file_to_be_uploaded, hide_progress=hide_progress)
+
+        self.clean_up()
+
+    def delete(self) -> None:
+        self._keep_last_n = 0
+        self._delete_old_remote_versions()
+
+    def _find_instances(self) -> None:
+        if self._cache_only_mode:
+            self._instances = self._find_instances_from_cache()
+        else:
+            self._instances = self._find_remote_instances()
+
+        self._instances = sorted(self._instances, key=lambda i: i.version)
+
+    def _find_instances_from_cache(self) -> List[DataInstance]:
+        logger.info(f"Fetching cached versions of {self._name}")
+
+        candidates = [
+            DataInstance(
+                name=CACHE_NAME_VERSION_SEPARATOR.join(
+                    f.name.split(CACHE_NAME_VERSION_SEPARATOR)[:-1]
+                ),
+                version=int(f.name.split(CACHE_NAME_VERSION_SEPARATOR)[-1]),
+                remote_path=f,
+                origin="filesystem",
+            )
+            for f in self.cache_path.glob(
+                f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}*"
+            )
+        ]
+
+        return [c for c in candidates if c.name == self._name]
+
+    def _check_mode_and_set_version(self) -> None:
+        if "+" in self._mode:
+            raise ValueError(
+                f"File mode `{self._mode}` is not allowed3, remove the `+`."
+            )
+
+        if "w" in self._mode:
+            if self._version is not None:
+                raise ValueError("Providing a version is not allowed in write mode.")
+
+            self._version = self._instances[-1].version + 1 if self._instances else 0
+
+        elif "r" in self._mode:
+            if not self._instances:
+                raise FileNotFoundError(
+                    f"File {self._name} not found. No versions are available."
+                )
+
+            if self._version is None:
+                self._version = self._instances[-1].version
+                logger.info(
+                    f"Latest version of {self._name} is {self._version} "
+                    + f"(from versions: {self.versions_pretty})"
+                )
+            elif self._version not in [i.version for i in self._instances]:
+                raise FileNotFoundError(
+                    f"File {self._name} not found with version {self._version}. "
+                    + f"(from versions: {self.versions_pretty})"
+                )
+        else:
+            raise ValueError("Unsupported file mode.")
+
+    @property
+    def versions_pretty(self) -> str:
+        return ", ".join((str(i.version) for i in self._instances))
+
+    def clean_up(self) -> None:
+        self._delete_old_remote_versions()
+        self._prune_cache()
+
+    def _prune_cache(self) -> None:
+        self.cache_path.mkdir(parents=True, exist_ok=True)
+
+        if self.max_cache_size is None:
+            return
+
+        allowed_size = human_readable_to_byte(self.max_cache_size)
+        assert allowed_size >= 0
+
+        least_recently_read = sorted(
+            [f for f in self.cache_path.glob("*")], key=lambda f: f.stat().st_atime
+        )
+
+        while sum(os.path.getsize(f) for f in least_recently_read) > allowed_size:
+            file = least_recently_read.pop(0)
+            logger.info(
+                f"Deleting file from cache to meet quota (max_cache_size={self.max_cache_size}): {file}"
+            )
+            os.unlink(file)
+
+    @abstractmethod
+    def _find_remote_instances(self) -> List[DataInstance]:
+        pass
+
+    @abstractmethod
+    def _download(
+        self, remote_path: Any, local_path: Path, hide_progress: bool
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def _upload(self, local_path: Path, hide_progress: bool) -> None:
+        pass
+
+    @abstractmethod
+    def _delete_old_remote_versions(self) -> None:
+        pass

great_ai/src/great_ai/large_file/large_file/large_file_local.py

@@ -0,0 +1,58 @@
+from pathlib import Path
+from typing import Any, List, Optional
+
+from great_ai.utilities.logger import get_logger
+
+from ..models import DataInstance
+from .large_file import LargeFile
+
+logger = get_logger("large_file")
+
+
+class LargeFileLocal(LargeFile):
+    def __init__(
+        self,
+        name: str,
+        mode: str = "r",
+        *,
+        buffering: int = -1,
+        encoding: Optional[str] = None,
+        errors: Optional[str] = None,
+        newline: Optional[str] = None,
+        version: Optional[int] = None,
+        keep_last_n: Optional[int] = None,
+    ):
+        super().__init__(
+            name,
+            mode,
+            buffering=buffering,
+            encoding=encoding,
+            errors=errors,
+            newline=newline,
+            version=version,
+            keep_last_n=keep_last_n,
+            cache_only_mode=True,
+        )
+
+    def _find_remote_instances(self) -> List[DataInstance]:
+        return []
+
+    def _download(
+        self, remote_path: Any, local_path: Path, hide_progress: bool
+    ) -> None:
+        raise NotImplementedError()
+
+    def _upload(self, local_path: Path, hide_progress: bool) -> None:
+        pass  # the "upload" is already done py the parent's caching mechanism
+
+    def _delete_old_remote_versions(self) -> None:
+        if self._keep_last_n is not None:
+            for i in (
+                self._instances[: -self._keep_last_n]
+                if self._keep_last_n > 0
+                else self._instances
+            ):
+                logger.info(
+                    f"Removing old version (keep_last_n={self._keep_last_n}): {i.remote_path}"
+                )
+                i.remote_path.unlink()

great_ai/src/great_ai/large_file/large_file/large_file_s3.py

@@ -0,0 +1,169 @@
+import configparser
+from functools import cached_property
+from pathlib import Path
+from typing import Any, List, Mapping, Optional, Union
+
+import boto3
+
+from great_ai.utilities.logger import get_logger
+
+from ..helper import DownloadProgressBar, UploadProgressBar
+from ..models import DataInstance
+from .large_file import LargeFile
+
+logger = get_logger("large_file")
+
+
+S3_NAME_VERSION_SEPARATOR = "/"
+
+
+class LargeFileS3(LargeFile):
+    """
+    Store large files in S3. Use local cache for speed up.
+
+    Examples:
+
+    ```
+    with LargeFile("test.txt", "w", keep_last_n=3) as f:
+        for i in range(1000000):
+            f.write('test\n')
+
+    with LargeFile("test.txt", "r") as f:
+        print(f.readlines()[0])
+
+    path_to_cached_text_file = LargeFile("test.txt", version=0).get()
+    ```
+
+    By default, files are stored in the ".cache" folder and the
+    least recently use is deleted after the overall size reaches 30 GBs.
+
+    Change it with the following properties.
+
+    ```
+    LargeFile.cache_path = Path(".cache")
+    LargeFile.max_cache_size = "30GB"
+    ```
+    """
+
+    region_name = None
+    access_key_id = None
+    secret_access_key = None
+    bucket_name = None
+    endpoint_url = None
+
+    @classmethod
+    def configure_credentials_from_file(
+        cls,
+        secrets_path: Union[Path, str],
+    ) -> None:
+        if isinstance(secrets_path, str):
+            secrets_path = Path(secrets_path)
+
+        if not secrets_path.exists():
+            raise FileNotFoundError(secrets_path.resolve())
+
+        credentials = configparser.ConfigParser()
+        credentials.read(secrets_path)
+        credentials.default_section
+        cls.configure_credentials(**credentials[credentials.default_section])
+
+    @classmethod
+    def configure_credentials(
+        cls,
+        *,
+        aws_region_name: str,
+        aws_access_key_id: str,
+        aws_secret_access_key: str,
+        large_files_bucket_name: str,
+        aws_endpoint_url: Optional[str] = None,
+        **_: Mapping[str, Any],
+    ) -> None:
+        cls.region_name = aws_region_name
+        cls.access_key_id = aws_access_key_id
+        cls.secret_access_key = aws_secret_access_key
+        cls.bucket_name = large_files_bucket_name
+        cls.endpoint_url = aws_endpoint_url
+
+    @cached_property
+    def _client(self) -> boto3.client:
+        if (
+            self.region_name is None
+            or self.access_key_id is None
+            or self.secret_access_key is None
+            or self.bucket_name is None
+        ):
+            raise ValueError(
+                "Please configure the S3 access options by calling LargeFile.configure_credentials or set offline_mode=True in the constructor."
+            )
+
+        return boto3.client(
+            "s3",
+            aws_access_key_id=self.access_key_id,
+            aws_secret_access_key=self.secret_access_key,
+            region_name=self.region_name,
+            endpoint_url=self.endpoint_url,
+        )
+
+    def _find_remote_instances(self) -> List[DataInstance]:
+        logger.debug(f"Fetching S3 versions of {self._name}")
+
+        found_objects = self._client.list_objects_v2(
+            Bucket=self.bucket_name, Prefix=self._name
+        )
+        return (
+            [
+                DataInstance(
+                    name=o["Key"].split(S3_NAME_VERSION_SEPARATOR)[0],
+                    version=int(o["Key"].split(S3_NAME_VERSION_SEPARATOR)[-1]),
+                    remote_path=o["Key"],
+                    origin="s3",
+                )
+                for o in found_objects["Contents"]
+                if o["Key"].split(S3_NAME_VERSION_SEPARATOR)[0] == self._name
+            ]
+            if "Contents" in found_objects
+            else []
+        )
+
+    def _download(
+        self, remote_path: Any, local_path: Path, hide_progress: bool
+    ) -> None:
+        logger.info(f"Downloading {remote_path} from S3")
+
+        size = self._client.head_object(Bucket=self.bucket_name, Key=remote_path)[
+            "ContentLength"
+        ]
+
+        self._client.download_file(
+            Bucket=self.bucket_name,
+            Key=hide_progress,
+            Filename=str(local_path),
+            Callback=None
+            if hide_progress
+            else DownloadProgressBar(name=str(remote_path), size=size, logger=logger),
+        )
+
+    def _upload(self, local_path: Path, hide_progress: bool) -> None:
+        key = f"{self._name}/{self.version}"
+        logger.info(f"Uploading {self._local_name} to S3 as {key}")
+
+        self._client.upload_file(
+            Filename=str(local_path),
+            Bucket=self.bucket_name,
+            Key=key,
+            Callback=None
+            if hide_progress
+            else UploadProgressBar(path=local_path, logger=logger),
+        )
+
+    def _delete_old_remote_versions(self) -> None:
+        if self._keep_last_n is not None:
+            for i in (
+                self._instances[: -self._keep_last_n]
+                if self._keep_last_n > 0
+                else self._instances
+            ):
+                logger.info(
+                    f"Removing old version from S3 (keep_last_n={self._keep_last_n}): {i.remote_path}"
+                )
+                self._client.delete_object(Bucket=self.bucket_name, Key=i.remote_path)

great_ai/src/great_ai/large_file/models/__init__.py

@@ -0,0 +1 @@
+from .data_instance import DataInstance

great_ai/src/great_ai/large_file/models/data_instance.py

@@ -0,0 +1,10 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+
+class DataInstance(BaseModel):
+    name: str
+    version: int
+    remote_path: Any
+    origin: Literal["filesystem", "mongodb", "s3"]

great_ai/src/great_ai/large_file/parse_arguments.py

@@ -7,12 +7,19 @@ def parse_arguments() -> Tuple[ArgumentParser, Namespace]:
         description="Store and version large files in S3; open them like regular files. Caching included.",
     )
 
+    parser.add_argument(
+        "--driver",
+        type=str,
+        help="choose which backend to use, available options: `local`, `s3`, `mongodb`",
+        required=True,
+    )
+
     parser.add_argument(
         "-s",
         "--secrets",
         type=str,
-        help="path to an .ini configration file with your S3 credentials",
+        help="path to an .ini configuration file with your S3 credentials",
-        required=True,
+        required=False,
     )
 
     parser.add_argument(

great_ai/src/great_ai/open_s3/__init__.py

@@ -1 +0,0 @@
-from .large_file import LargeFile

great_ai/src/great_ai/open_s3/large_file.py

@@ -1,373 +0,0 @@
-import configparser
-import os
-import shutil
-import tempfile
-from pathlib import Path
-from types import TracebackType
-from typing import IO, Any, List, Mapping, Optional, Type, Union, cast
-
-import boto3
-
-from great_ai.utilities.logger import get_logger
-
-from .helper import DownloadProgressBar, UploadProgressBar, human_readable_to_byte
-
-logger = get_logger("open_s3")
-
-
-class LargeFile:
-    """
-    Store large files in S3. Use local cache for speed up.
-
-    Examples:
-
-    ```
-    with LargeFile("test.txt", "w", keep_last_n=3) as f:
-        for i in range(1000000):
-            f.write('test\n')
-
-    with LargeFile("test.txt", "r") as f:
-        print(f.readlines()[0])
-
-    path_to_cached_text_file = LargeFile("test.txt", version=0).get()
-    ```
-
-    By default, files are stored in the ".cache" folder and the
-    least recently use is deleted after the overall size reaches 30 GBs.
-
-    Change it with the following properties.
-
-    ```
-    LargeFile.cache_path = Path(".cache")
-    LargeFile.max_cache_size = "30GB"
-    ```
-    """
-
-    region_name = None
-    access_key_id = None
-    secret_access_key = None
-    bucket_name = None
-    endpoint_url = None
-
-    cache_path = Path(".cache")
-    max_cache_size: Optional[str] = "30GB"
-
-    def __init__(
-        self,
-        name: str,
-        mode: str = "r",
-        *,
-        buffering: int = -1,
-        encoding: Optional[str] = None,
-        errors: Optional[str] = None,
-        newline: Optional[str] = None,
-        version: Optional[int] = None,
-        keep_last_n: Optional[int] = None,
-        offline_mode: bool = False,
-    ):
-        self._name: str = name
-        self._version: int = cast(int, version)
-        self._mode: str = mode
-        self._keep_last_n = keep_last_n
-        self._offline_mode = offline_mode
-
-        self._buffering = buffering
-        self._encoding = encoding
-        self._errors = errors
-        self._newline = newline
-
-        LargeFile.cache_path.mkdir(parents=True, exist_ok=True)
-
-        self._find_versions()
-        self._check_mode_and_set_version()
-
-    @classmethod
-    def configure_credentials_from_file(
-        cls,
-        secrets_path: Union[Path, str],
-    ) -> None:
-        if isinstance(secrets_path, str):
-            secrets_path = Path(secrets_path)
-
-        if not secrets_path.exists():
-            raise FileNotFoundError(secrets_path.resolve())
-
-        credentials = configparser.ConfigParser()
-        credentials.read(secrets_path)
-        credentials.default_section
-        cls.configure_credentials(**credentials[credentials.default_section])
-
-    @classmethod
-    def configure_credentials(
-        cls,
-        *,
-        aws_region_name: str,
-        aws_access_key_id: str,
-        aws_secret_access_key: str,
-        large_files_bucket_name: str,
-        aws_endpoint_url: Optional[str] = None,
-        **_: Mapping[str, Any],
-    ) -> None:
-        cls.region_name = aws_region_name
-        cls.access_key_id = aws_access_key_id
-        cls.secret_access_key = aws_secret_access_key
-        cls.bucket_name = large_files_bucket_name
-        cls.endpoint_url = aws_endpoint_url
-
-    def __enter__(self) -> IO:
-        self._file: IO[Any] = (
-            tempfile.NamedTemporaryFile(
-                mode=self._mode,
-                buffering=self._buffering,
-                encoding=self._encoding,
-                newline=self._newline,
-                errors=self._errors,
-                delete=False,
-                prefix="large-file-",
-            )
-            if "w" in self._mode
-            else open(
-                self.get(),
-                mode=self._mode,
-                buffering=self._buffering,
-                encoding=self._encoding,
-                newline=self._newline,
-                errors=self._errors,
-            )
-        )
-
-        return self._file
-
-    def __exit__(
-        self,
-        type: Optional[Type[BaseException]],
-        exc: Optional[BaseException],
-        traceback: Optional[TracebackType],
-    ) -> bool:
-        self._file.close()
-
-        if type is None:
-            if "w" in self._mode:
-                self.push(Path(self._file.name))
-                os.unlink(self._file.name)
-        else:
-            logger.exception("Could not finish operation.")
-
-        return True
-
-    @property
-    def version_ids(self) -> List[int]:
-        return [self._get_version_from_key(key) for key in self._versions]
-
-    @property
-    def version(self) -> int:
-        return self._version
-
-    def get(self, hide_progress: bool = False) -> Path:
-        key = next(
-            key
-            for key in self._versions
-            if self._get_version_from_key(key) == self._version
-        )
-
-        destination = self.cache_path / self._local_name
-        if not destination.exists():
-            logger.info(
-                f"File {self._local_name} does not exist locally, starting download from S3"
-            )
-
-            with tempfile.TemporaryDirectory() as tmp:
-                tmp_file_archive = Path(tmp) / f"{self._local_name}.tar.gz"
-
-                size = self._client.head_object(Bucket=self.bucket_name, Key=key)[
-                    "ContentLength"
-                ]
-                self._client.download_file(
-                    Bucket=self.bucket_name,
-                    Key=key,
-                    Filename=str(tmp_file_archive),
-                    Callback=None
-                    if hide_progress
-                    else DownloadProgressBar(name=str(key), size=size, logger=logger),
-                )
-                logger.info(f"Decompressing {self._local_name}")
-                shutil.unpack_archive(str(tmp_file_archive), tmp, "gztar")
-                tmp_file = Path(tmp) / self._local_name
-                shutil.move(str(tmp_file), str(destination))
-        else:
-            logger.info(f"File {self._local_name} found in cache")
-
-        return destination
-
-    def push(self, path: Union[Path, str], hide_progress: bool = False) -> None:
-        if isinstance(path, str):
-            path = Path(path)
-
-        with tempfile.TemporaryDirectory() as tmp:
-            if path.is_file():
-                logger.info(f"Copying file for {self._local_name}")
-                copy: Any = shutil.copy
-            else:
-                logger.info(f"Copying directory for {self._local_name}")
-                copy = shutil.copytree
-
-            try:
-                # Make local copy in the cache
-                copy(str(path), str(self.cache_path / self._local_name))
-            except shutil.SameFileError:
-                pass  # No worries
-
-            copy(str(path), str(Path(tmp) / self._local_name))
-
-            with tempfile.TemporaryDirectory() as tmp2:
-                # A directory has to be zipped and it cannot contain the output of the zipping
-                logger.info(f"Compressing {self._local_name}")
-                shutil.make_archive(
-                    str(Path(tmp2) / self._local_name),
-                    "gztar",
-                    tmp,
-                )
-
-                logger.info(f"Uploading {self._local_name} to S3 from {path}")
-
-                file_to_be_uploaded = Path(tmp2) / f"{self._local_name}.tar.gz"
-                self._client.upload_file(
-                    Filename=str(file_to_be_uploaded),
-                    Bucket=self.bucket_name,
-                    Key=self._s3_name,
-                    Callback=None
-                    if hide_progress
-                    else UploadProgressBar(path=file_to_be_uploaded, logger=logger),
-                )
-
-        self.clean_up()
-
-    def delete(self) -> None:
-        self._keep_last_n = 0
-        self._delete_old_versions_from_s3()
-
-    def clean_up(self) -> None:
-        self._delete_old_versions_from_s3()
-        self._delete_old_versions_from_disk()
-
-    def _create_client(self) -> None:
-        if (
-            self.region_name is None
-            or self.access_key_id is None
-            or self.secret_access_key is None
-            or self.bucket_name is None
-        ):
-            raise ValueError(
-                "Please configure the S3 access options by calling LargeFile.configure_credentials or set offline_mode=True in the constructor."
-            )
-
-        self._client = boto3.client(
-            "s3",
-            aws_access_key_id=self.access_key_id,
-            aws_secret_access_key=self.secret_access_key,
-            region_name=self.region_name,
-            endpoint_url=self.endpoint_url,
-        )
-
-    def _find_versions(self) -> None:
-        if self._offline_mode:
-            versions: Union[List[str], List[Path]] = self._fetch_versions_from_cache()
-        else:
-            self._create_client()
-            versions = self._fetch_versions_from_s3()
-
-        self._versions = sorted(versions, key=self._get_version_from_key)
-
-    def _fetch_versions_from_cache(self) -> List[Path]:
-        logger.debug(f"Fetching offline versions of {self._name}")
-
-        return list(self.cache_path.glob(f"{self._local_name}-*"))
-
-    def _fetch_versions_from_s3(self) -> List[str]:
-        logger.debug(f"Fetching online versions of {self._name}")
-
-        found_objects = self._client.list_objects_v2(
-            Bucket=self.bucket_name, Prefix=self._name
-        )
-        return (
-            [o["Key"] for o in found_objects["Contents"]]
-            if "Contents" in found_objects
-            else []
-        )
-
-    def _check_mode_and_set_version(self) -> None:
-        if "+" in self._mode:
-            raise ValueError("Read-write mode is not allowed.")
-
-        if "w" in self._mode:
-            if self._version is not None:
-                raise ValueError("Providing a version is not allowed in write mode.")
-
-            self._version = self.version_ids[-1] + 1 if self.version_ids else 0
-
-        elif "r" in self._mode:
-            if not self.version_ids:
-                raise FileNotFoundError(
-                    f"File {self._name} not found. No versions are available."
-                )
-
-            if self._version is None:
-                self._version = self.version_ids[-1]
-                logger.info(
-                    f"Latest version of {self._name} is {self._version} "
-                    + f"(from versions: {', '.join((str(v) for v in self.version_ids))})"
-                )
-
-            elif self._version not in self.version_ids:
-                raise FileNotFoundError(
-                    f"File {self._name} not found with version {self._version}. "
-                    + f"(from versions: {', '.join((str(v) for v in self.version_ids))})"
-                )
-        else:
-            raise ValueError("Unsupported file mode.")
-
-    @property
-    def _local_name(self) -> str:
-        return f"{self._name}-{self._version}"
-
-    @property
-    def _s3_name(self) -> str:
-        return f"{self._name}/{self._version}"
-
-    @staticmethod
-    def _get_version_from_key(key: Union[str, Path]) -> int:
-        if isinstance(key, Path):
-            return int(key.name.split("-")[-1])
-        return int(key.split("/")[-1])
-
-    def _delete_old_versions_from_s3(self) -> None:
-        if self._keep_last_n is not None:
-            for key in (
-                self._versions[: -self._keep_last_n]
-                if self._keep_last_n > 0
-                else self._versions
-            ):
-                logger.info(
-                    f"Removing old version (keep_last_n={self._keep_last_n}): {key}"
-                )
-                self._client.delete_object(Bucket=self.bucket_name, Key=key)
-
-    def _delete_old_versions_from_disk(self) -> None:
-        self.cache_path.mkdir(parents=True, exist_ok=True)
-
-        if self.max_cache_size is None:
-            return
-
-        allowed_size = human_readable_to_byte(self.max_cache_size)
-        assert allowed_size >= 0
-
-        least_recently_read = sorted(
-            [f for f in self.cache_path.glob("*")], key=lambda f: f.stat().st_atime
-        )
-
-        while sum(os.path.getsize(f) for f in least_recently_read) > allowed_size:
-            file = least_recently_read.pop(0)
-            logger.info(
-                f"Deleting file from cache to meet quota (max_cache_size={self.max_cache_size}): {file}"
-            )
-            os.unlink(file)

great_ai/tests/large_file/test_human_readable_to_byte.py

@@ -1,6 +1,6 @@
 import unittest
 
-from src.great_ai.open_s3.helper import human_readable_to_byte
+from src.great_ai.large_file.helper import human_readable_to_byte
 
 
 class TestHumanReadableToByte(unittest.TestCase):

great_ai/tests/large_file/test_large_file.py

@@ -8,7 +8,7 @@ import boto3
 PATH = Path(__file__).parent.resolve()
 
 
-from src.great_ai import LargeFile
+from src.great_ai import LargeFileS3
 
 credentials = {
     "aws_region_name": "your_region_like_eu-west-2",
@@ -19,15 +19,15 @@ credentials = {
 }
 
 
-class TestLargeFile(unittest.TestCase):
+class TestLargeFileS3(unittest.TestCase):
     def test_uninitialized(self) -> None:
-        self.assertRaises(ValueError, LargeFile, "test-file")
+        self.assertRaises(ValueError, LargeFileS3, "test-file")
 
     def test_bad_file_modes(self) -> None:
-        self.assertRaises(ValueError, LargeFile, "test-file", "w", version=3)
+        self.assertRaises(ValueError, LargeFileS3, "test-file", "w", version=3)
-        self.assertRaises(ValueError, LargeFile, "test-file", "wb", version=3)
+        self.assertRaises(ValueError, LargeFileS3, "test-file", "wb", version=3)
-        self.assertRaises(ValueError, LargeFile, "test-file", "w+r")
+        self.assertRaises(ValueError, LargeFileS3, "test-file", "w+r")
-        self.assertRaises(ValueError, LargeFile, "test-file", "test")
+        self.assertRaises(ValueError, LargeFileS3, "test-file", "test")
 
     @patch.object(boto3, "client")
     def test_initialized_with_dict(self, client: Any) -> None:
@@ -52,14 +52,14 @@ class TestLargeFile(unittest.TestCase):
         )
         boto3.client = Mock(return_value=s3)
 
-        LargeFile.configure_credentials(
+        LargeFileS3.configure_credentials(
             aws_region_name=credentials["aws_region_name"],
             aws_access_key_id=credentials["aws_access_key_id"],
             aws_secret_access_key=credentials["aws_secret_access_key"],
             large_files_bucket_name=credentials["large_files_bucket_name"],
             aws_endpoint_url=credentials["aws_endpoint_url"],
         )
-        lf = LargeFile("test-file")
+        lf = LargeFileS3("test-file")
 
         boto3.client.assert_called_once_with(
             "s3",
@@ -75,7 +75,6 @@ class TestLargeFile(unittest.TestCase):
 
         self.assertEqual(lf._version, 2)
         self.assertEqual(lf._local_name, "test-file-2")
-        self.assertEqual(lf._s3_name, "test-file/2")
 
     @patch.object(boto3, "client")
     def test_initialized_with_file(self, client: Any) -> None:
@@ -101,8 +100,8 @@ class TestLargeFile(unittest.TestCase):
 
         boto3.client = Mock(return_value=s3)
 
-        LargeFile.configure_credentials_from_file(PATH / "../../example_secrets.ini")
+        LargeFileS3.configure_credentials_from_file(PATH / "../../example_secrets.ini")
-        lf = LargeFile("test-file")
+        lf = LargeFileS3("test-file")
 
         boto3.client.assert_called_once_with(
             "s3",
@@ -119,4 +118,3 @@ class TestLargeFile(unittest.TestCase):
 
         self.assertEqual(lf._version, 2)
         self.assertEqual(lf._local_name, "test-file-2")
-        self.assertEqual(lf._s3_name, "test-file/2")