UncountablePythonSDK 0.0.24__py3-none-any.whl → 0.0.131__py3-none-any.whl

pkgs/argument_parser/case_convert.py

@@ -4,9 +4,10 @@ import re
 
 @functools.lru_cache(maxsize=500000)
 def snake_to_camel_case(o: str) -> str:
-    return "".join([
-        part.title() if i > 0 else part for i, part in enumerate(o.split("_"))
-    ])
+    tokens = o.split("_")
+    if len(tokens) < 2:
+        return o
+    return "".join([part.title() if i > 0 else part for i, part in enumerate(tokens)])
 
 
 def kebab_to_pascal_case(o: str) -> str:

pkgs/filesystem_utils/__init__.py

@@ -0,0 +1,20 @@
+from ._gdrive_session import GDriveSession as GDriveSession
+from ._gdrive_session import delete_gdrive_file as delete_gdrive_file
+from ._gdrive_session import download_gdrive_file as download_gdrive_file
+from ._gdrive_session import list_gdrive_files as list_gdrive_files
+from ._gdrive_session import move_gdrive_file as move_gdrive_file
+from ._gdrive_session import upload_file_gdrive as upload_file_gdrive
+from ._local_session import LocalSession as LocalSession
+from ._s3_session import S3Session as S3Session
+from ._sftp_session import SFTPSession as SFTPSession
+from ._sftp_session import list_sftp_files as list_sftp_files
+from ._sftp_session import move_sftp_files as move_sftp_files
+from .file_type_utils import FileObjectData as FileObjectData
+from .file_type_utils import FileSystemFileReference as FileSystemFileReference
+from .file_type_utils import FileSystemObject as FileSystemObject
+from .file_type_utils import FileSystemS3Config as FileSystemS3Config
+from .file_type_utils import FileSystemSFTPConfig as FileSystemSFTPConfig
+from .file_type_utils import FileTransfer as FileTransfer
+from .file_type_utils import IncompatibleFileReference as IncompatibleFileReference
+from .file_type_utils import RemoteObjectReference as RemoteObjectReference
+from .filesystem_session import FileSystemSession as FileSystemSession

pkgs/filesystem_utils/_blob_session.py

@@ -0,0 +1,144 @@
+from io import BytesIO
+
+from azure.core.credentials import AzureSasCredential
+from azure.storage.blob import BlobServiceClient, ContainerClient
+
+from pkgs.filesystem_utils.file_type_utils import (
+    FileObjectData,
+    FileSystemBlobConfig,
+    FileSystemFileReference,
+    FileSystemObject,
+    FileTransfer,
+    IncompatibleFileReference,
+)
+
+from .filesystem_session import FileSystemSession
+
+
+def _add_slash(prefix: str) -> str:
+    if len(prefix) > 0 and prefix[-1] != "/":
+        prefix = prefix + "/"
+    return prefix
+
+
+class BlobSession(FileSystemSession):
+    config: FileSystemBlobConfig
+
+    def __init__(self, blob_config: FileSystemBlobConfig) -> None:
+        super().__init__()
+        self.config = blob_config
+
+    def start(self) -> None:
+        self.service_client: BlobServiceClient | None = BlobServiceClient(
+            self.config.account_url, credential=self.config.credential
+        )
+        self.container_client: ContainerClient | None = (
+            self.service_client.get_container_client(self.config.container)
+        )
+
+    def __enter__(self) -> "BlobSession":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        self.service_client = None
+        self.container_client = None
+
+    def list_files(
+        self,
+        dir_path: FileSystemObject,
+        *,
+        recursive: bool = False,
+        valid_extensions: list[str] | None = None,
+    ) -> list[FileSystemObject]:
+        if not isinstance(dir_path, FileSystemFileReference):
+            raise IncompatibleFileReference()
+
+        assert self.service_client is not None and self.container_client is not None, (
+            "call to list_files on uninitialized blob session"
+        )
+
+        filesystem_file_references: list[FileSystemObject] = []
+        prefix = _add_slash(dir_path.filepath)
+        for blob in self.container_client.list_blobs(name_starts_with=prefix):
+            if not recursive and (
+                blob.name == prefix or "/" in blob.name[len(prefix) :]
+            ):
+                continue
+            if valid_extensions is None or any(
+                blob.name.endswith(valid_extension)
+                for valid_extension in valid_extensions
+            ):
+                filesystem_file_references.append(
+                    FileSystemFileReference(
+                        filepath=blob.name,
+                    )
+                )
+
+        return filesystem_file_references
+
+    def download_files(
+        self,
+        filepaths: list[FileSystemObject],
+    ) -> list[FileObjectData]:
+        downloaded_files: list[FileObjectData] = []
+        assert self.service_client is not None and self.container_client is not None, (
+            "call to download_files on uninitialized blob session"
+        )
+
+        for file_object in filepaths:
+            if (
+                not isinstance(file_object, FileSystemFileReference)
+                or file_object.filename is None
+            ):
+                raise IncompatibleFileReference()
+
+            blob_client = self.container_client.get_blob_client(file_object.filepath)
+            download_stream = blob_client.download_blob()
+            file_data = download_stream.readall()
+            downloaded_files.append(
+                FileObjectData(
+                    file_data=file_data,
+                    file_IO=BytesIO(file_data),
+                    filename=file_object.filename,
+                    filepath=file_object.filepath,
+                )
+            )
+
+        return downloaded_files
+
+    def move_files(self, file_mappings: list[FileTransfer]) -> None:
+        assert self.service_client is not None and self.container_client is not None, (
+            "call to move_files on uninitialized blob session"
+        )
+
+        for src_file, dest_file in file_mappings:
+            if not isinstance(src_file, FileSystemFileReference) or not isinstance(
+                dest_file, FileSystemFileReference
+            ):
+                raise IncompatibleFileReference()
+
+            source_blob_client = self.container_client.get_blob_client(
+                src_file.filepath
+            )
+            dest_blob_client = self.container_client.get_blob_client(dest_file.filepath)
+
+            source_url = (
+                f"{source_blob_client.url}?{self.config.credential.signature}"
+                if isinstance(self.config.credential, AzureSasCredential)
+                else source_blob_client.url
+            )
+
+            dest_blob_client.start_copy_from_url(source_url)
+            source_blob_client.delete_blob()
+
+    def delete_files(self, filepaths: list[FileSystemObject]) -> None:
+        assert self.service_client is not None and self.container_client is not None, (
+            "call to delete_files on uninitialized blob session"
+        )
+        for file_object in filepaths:
+            if not isinstance(file_object, FileSystemFileReference):
+                raise IncompatibleFileReference()
+
+            blob_client = self.container_client.get_blob_client(file_object.filepath)
+            blob_client.delete_blob()

pkgs/filesystem_utils/_gdrive_session.py

@@ -0,0 +1,309 @@
+import os
+from io import BytesIO
+from typing import Any
+
+from google.oauth2 import service_account
+from googleapiclient.discovery import build as build_gdrive_connection
+from googleapiclient.errors import HttpError
+from googleapiclient.http import MediaIoBaseDownload, MediaIoBaseUpload
+from tqdm import tqdm
+
+from pkgs.filesystem_utils.file_type_utils import (
+    FileObjectData,
+    FileSystemFileReference,
+    FileSystemObject,
+    FileTransfer,
+    IncompatibleFileReference,
+    RemoteObjectReference,
+)
+
+from .filesystem_session import FileSystemSession
+
+# NOTE: google apis do not have static types
+GDriveResource = Any
+
+
+def download_gdrive_file(
+    gdrive_connection: GDriveResource,
+    file_id: str,
+    filename: str,
+    mime_type: str,
+    *,
+    verbose: bool = False,
+) -> FileObjectData | None:
+    if "folder" in mime_type:
+        if verbose:
+            print(f"{filename} is a folder and will not be downloaded.")
+        return None
+    elif "google-apps" in mime_type:
+        # Handle google workspace doc
+        if "spreadsheet" in mime_type:
+            if verbose:
+                print(f"{filename} is a Google Sheet, exporting.")
+            file_request = gdrive_connection.files().export_media(
+                fileId=file_id, mimeType="text/csv"
+            )
+            filename += ".csv"
+        elif "document" in mime_type:
+            if verbose:
+                print(f"{filename} is a Google Doc, exporting.")
+            file_request = gdrive_connection.files().export_media(
+                fileId=file_id, mimeType="application/msword"
+            )
+            filename += ".doc"
+        else:
+            if verbose:
+                print(f"{filename} is an unsupported google workspace filetype.")
+                print(f"Skipping. mimeType: {mime_type}.")
+            return None
+    else:
+        file_request = gdrive_connection.files().get_media(fileId=file_id)
+
+    file_handler = BytesIO()
+    downloader = MediaIoBaseDownload(file_handler, file_request)
+    download_complete = False
+    while not download_complete:
+        _status, download_complete = downloader.next_chunk()
+
+    file_handler.seek(0)
+    file_data = file_handler.read()
+    return FileObjectData(
+        file_data=file_data,
+        file_IO=BytesIO(file_data),
+        filename=filename,
+        filepath=file_id,
+        metadata={"id": file_id},
+        mime_type=mime_type,
+    )
+
+
+def list_gdrive_files(
+    gdrive_connection: GDriveResource, gdrive_folder_id: str, *, recurse: bool = False
+) -> list[dict[str, str]]:
+    query = f"parents = '{gdrive_folder_id}'"
+    print("Listing files", end="", flush=True)
+    paginated_files_in_folder = [
+        (
+            gdrive_connection.files()
+            .list(
+                q=query,
+                corpora="allDrives",
+                includeItemsFromAllDrives=True,
+                supportsAllDrives=True,
+            )
+            .execute()
+        )
+    ]
+    while paginated_files_in_folder[-1].get("nextPageToken") is not None:
+        print(".", end="", flush=True)
+        paginated_files_in_folder.append(
+            gdrive_connection.files()
+            .list(
+                q=query,
+                corpora="allDrives",
+                includeItemsFromAllDrives=True,
+                supportsAllDrives=True,
+                pageToken=paginated_files_in_folder[-1]["nextPageToken"],
+            )
+            .execute()
+        )
+    print()
+    # Get available files: https://developers.google.com/drive/api/v3/manage-downloads#python
+    files: list[dict[str, str]] = []
+    for files_in_folder in paginated_files_in_folder:
+        files.extend(files_in_folder.get("files", []))
+    subfiles: list[dict[str, str]] = []
+    if recurse:
+        for file in files:
+            if file["mimeType"] == "application/vnd.google-apps.folder":
+                subfiles.extend(
+                    list_gdrive_files(
+                        gdrive_connection=gdrive_connection,
+                        gdrive_folder_id=file["id"],
+                        recurse=True,
+                    )
+                )
+    return [*files, *subfiles]
+
+
+def upload_file_gdrive(
+    gdrive_connection: GDriveResource,
+    src_file: BytesIO,
+    mime_type: str,
+    dest_folder_id: str,
+    dest_filename: str,
+) -> None:
+    file_metadata = {"name": dest_filename, "parents": [dest_folder_id]}
+    media = MediaIoBaseUpload(src_file, mimetype=mime_type)
+    try:
+        gdrive_connection.files().create(
+            body=file_metadata, media_body=media, fields="id", supportsAllDrives=True
+        ).execute()
+    except HttpError:
+        print("FileSystemObject Upload to GDrive Unsuccessful")
+
+
+def move_gdrive_file(
+    gdrive_connection: GDriveResource,
+    src_file_id: str,
+    dest_folder_id: str,
+    *,
+    dest_filename: str | None = None,
+) -> None:
+    # Retrieve the existing parents to remove
+    file = (
+        gdrive_connection.files()
+        .get(fileId=src_file_id, fields="parents, name", supportsTeamDrives=True)
+        .execute()
+    )
+
+    new_filename = file["name"]
+    if dest_filename is not None:
+        new_filename = dest_filename
+    previous_parents = ",".join(file.get("parents"))
+    metadata = {"name": new_filename}
+    gdrive_connection.files().update(
+        fileId=src_file_id, body=metadata, fields="name", supportsTeamDrives=True
+    ).execute()
+    gdrive_connection.files().update(
+        fileId=src_file_id,
+        addParents=dest_folder_id,
+        removeParents=previous_parents,
+        fields="id, parents",
+        supportsTeamDrives=True,
+    ).execute()
+
+
+def delete_gdrive_file(gdrive_connection: GDriveResource, file_id: str) -> None:
+    gdrive_connection.files().delete(fileId=file_id, supportsAllDrives=True).execute()
+
+
+class GDriveSession(FileSystemSession):
+    def __init__(self, service_account_json_path: str) -> None:
+        super().__init__()
+        self.service_account_json_path = service_account_json_path
+
+    def start(self) -> None:
+        credentials = service_account.Credentials.from_service_account_file(  # type: ignore[no-untyped-call]
+            self.service_account_json_path
+        )
+        gdrive_connection = build_gdrive_connection(
+            "drive", "v3", credentials=credentials
+        )
+        self.connection = gdrive_connection
+
+    def list_files(
+        self,
+        dir_path: FileSystemObject,
+        *,
+        recursive: bool = False,
+        valid_file_extensions: tuple[str, ...] | None = None,
+    ) -> list[FileSystemObject]:
+        if not isinstance(dir_path, RemoteObjectReference):
+            raise IncompatibleFileReference(
+                "Incompatible FileSystemObject to GDriveSession.list_files"
+            )
+        if not dir_path.is_dir:
+            raise IncompatibleFileReference(
+                "FileSystemObject does not reference a directory"
+            )
+        files = list_gdrive_files(self.connection, dir_path.file_id, recurse=recursive)
+        gdrive_files: list[FileSystemObject] = []
+        for file_context in files:
+            if (
+                valid_file_extensions is not None
+                and os.path.splitext(file_context["name"])[1]
+                not in valid_file_extensions
+            ):
+                continue
+            gdrive_files.append(
+                RemoteObjectReference(
+                    file_id=file_context["id"],
+                    mime_type=file_context["mimeType"],
+                    filename=file_context["name"],
+                )
+            )
+        return gdrive_files
+
+    def delete_files(self, filepaths: list[FileSystemObject]) -> None:
+        """Warning:
+        Security account must have sufficient permissions to perform delete!
+        https://developers.google.com/drive/api/v3/reference/files/delete?hl=en
+        https://developers.google.com/drive/api/v3/ref-roles
+        """
+        for file_object in filepaths:
+            if not isinstance(file_object, RemoteObjectReference):
+                raise IncompatibleFileReference(
+                    "Incompatible FileSystemObject provided to GDriveSession.delete_files"
+                )
+            delete_gdrive_file(self.connection, file_object.file_id)
+
+    def move_files(self, file_mappings: list[FileTransfer]) -> None:
+        for src_file, dest_file in file_mappings:
+            if (
+                isinstance(src_file, FileSystemFileReference)
+                or not isinstance(dest_file, RemoteObjectReference)
+                or not dest_file.is_dir
+                or (isinstance(src_file, RemoteObjectReference) and src_file.is_dir)
+            ):
+                continue
+            new_filename = dest_file.filename
+            if isinstance(src_file, RemoteObjectReference):
+                if new_filename is not None:
+                    move_gdrive_file(
+                        self.connection,
+                        src_file.file_id,
+                        dest_file.file_id,
+                        dest_filename=new_filename,
+                    )
+                else:
+                    move_gdrive_file(
+                        self.connection, src_file.file_id, dest_file.file_id
+                    )
+            elif isinstance(src_file, FileObjectData):
+                if src_file.mime_type is None:
+                    raise IncompatibleFileReference(
+                        "No mime_type present on source file data."
+                    )
+                new_filename = src_file.filename
+                if dest_file.filename is not None:
+                    new_filename = dest_file.filename
+                upload_file_gdrive(
+                    self.connection,
+                    src_file.file_IO,
+                    src_file.mime_type,
+                    dest_file.file_id,
+                    new_filename,
+                )
+            else:
+                raise IncompatibleFileReference(
+                    "Unrecognized file reference in FileTransfer object"
+                )
+
+    def download_files(self, filepaths: list[FileSystemObject]) -> list[FileObjectData]:
+        downloaded_files: list[FileObjectData] = []
+        print(f"Downloading {len(filepaths)} files")
+        for file_object in tqdm(filepaths):
+            if (
+                not isinstance(file_object, RemoteObjectReference)
+                or file_object.filename is None
+            ):
+                raise IncompatibleFileReference(
+                    "Incompatible FileSystemObject included in filepaths"
+                )
+            downloaded_file = download_gdrive_file(
+                self.connection,
+                file_object.file_id,
+                file_object.filename,
+                file_object.mime_type,
+            )
+            if downloaded_file is not None:
+                downloaded_files.append(downloaded_file)
+        return downloaded_files
+
+    def __enter__(self) -> "GDriveSession":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        self.connection.close()

pkgs/filesystem_utils/_local_session.py

@@ -0,0 +1,69 @@
+import os
+from io import BytesIO
+
+from pkgs.filesystem_utils.file_type_utils import (
+    FileObjectData,
+    FileSystemFileReference,
+    FileSystemObject,
+    FileTransfer,
+    IncompatibleFileReference,
+)
+
+from .filesystem_session import FileSystemSession
+
+
+class LocalSession(FileSystemSession):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def start(self) -> None:
+        return None
+
+    def __enter__(self) -> "LocalSession":
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        return None
+
+    def move_files(self, file_mappings: list[FileTransfer]) -> None:
+        for src_file, dest_file in file_mappings:
+            if not (
+                isinstance(src_file, FileSystemFileReference)
+                and isinstance(dest_file, FileSystemFileReference)
+            ):
+                raise IncompatibleFileReference()
+            os.rename(src_file.filepath, dest_file.filepath)
+
+    def download_files(self, filepaths: list[FileSystemObject]) -> list[FileObjectData]:
+        downloaded_files: list[FileObjectData] = []
+        for file_object in filepaths:
+            if (
+                not isinstance(file_object, FileSystemFileReference)
+                or file_object.filename is None
+            ):
+                raise IncompatibleFileReference()
+            with open(file_object.filepath, "rb") as file_data:
+                file_bytes = file_data.read()
+            downloaded_files.append(
+                FileObjectData(
+                    file_bytes,
+                    BytesIO(file_bytes),
+                    file_object.filename,
+                    filepath=file_object.filepath,
+                )
+            )
+        return downloaded_files
+
+    def list_files(
+        self, dir_path: FileSystemObject, *, recursive: bool = False
+    ) -> list[FileSystemObject]:
+        if not isinstance(dir_path, FileSystemFileReference) or not os.path.isdir(
+            dir_path.filepath
+        ):
+            raise IncompatibleFileReference()
+        if recursive:
+            raise NotImplementedError("recursive not implemented for local session")
+        return [
+            FileSystemFileReference(os.path.join(dir_path.filepath, filename))
+            for filename in os.listdir(dir_path.filepath)
+        ]

pkgs/filesystem_utils/_s3_session.py

@@ -0,0 +1,118 @@
+from io import BytesIO
+
+from boto3.session import Session
+from mypy_boto3_s3.service_resource import Bucket
+
+from pkgs.filesystem_utils.file_type_utils import (
+    FileObjectData,
+    FileSystemFileReference,
+    FileSystemObject,
+    FileSystemS3Config,
+    FileTransfer,
+    IncompatibleFileReference,
+)
+
+from .filesystem_session import FileSystemSession
+
+
+def _add_slash(prefix: str) -> str:
+    if len(prefix) > 0 and prefix[-1] != "/":
+        prefix = prefix + "/"
+    return prefix
+
+
+class S3Session(FileSystemSession):
+    config: FileSystemS3Config
+
+    def __init__(self, s3_config: FileSystemS3Config) -> None:
+        super().__init__()
+        self.config = s3_config
+
+    def start(self) -> None:
+        session = Session(region_name=self.config.region_name)
+        s3_resource = session.resource(
+            service_name="s3",
+            endpoint_url=self.config.endpoint_url,
+            aws_access_key_id=self.config.access_key_id,
+            aws_secret_access_key=self.config.secret_access_key,
+            aws_session_token=self.config.session_token,
+        )
+
+        self.bucket: Bucket | None = s3_resource.Bucket(self.config.bucket_name)
+
+    def __enter__(self) -> "S3Session":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        self.bucket = None
+
+    def list_files(
+        self,
+        dir_path: FileSystemObject,
+        *,
+        recursive: bool = False,
+        valid_extensions: list[str] | None = None,
+    ) -> list[FileSystemObject]:
+        if not isinstance(dir_path, FileSystemFileReference):
+            raise IncompatibleFileReference()
+
+        assert self.bucket is not None, "call to list_files on uninitialized s3 session"
+
+        filesystem_references: list[FileSystemObject] = []
+        prefix = _add_slash(dir_path.filepath)
+        for obj in self.bucket.objects.filter(Prefix=prefix):
+            if not recursive and (obj.key == prefix or "/" in obj.key[len(prefix) :]):
+                continue
+            if valid_extensions is None or any(
+                obj.key.endswith(valid_extension)
+                for valid_extension in valid_extensions
+            ):
+                filesystem_references.append(FileSystemFileReference(obj.key))
+
+        return filesystem_references
+
+    def download_files(
+        self,
+        filepaths: list[FileSystemObject],
+    ) -> list[FileObjectData]:
+        downloaded_files: list[FileObjectData] = []
+        assert self.bucket is not None, (
+            "call to download_files on uninitialized s3 session"
+        )
+
+        for file_object in filepaths:
+            if (
+                not isinstance(file_object, FileSystemFileReference)
+                or file_object.filename is None
+            ):
+                raise IncompatibleFileReference()
+            s3_file_obj = self.bucket.Object(file_object.filepath)
+            response = s3_file_obj.get()
+            file_obj_bytes = response["Body"].read()
+            downloaded_files.append(
+                FileObjectData(
+                    file_data=file_obj_bytes,
+                    file_IO=BytesIO(file_obj_bytes),
+                    filename=file_object.filename,
+                    filepath=file_object.filepath,
+                )
+            )
+
+        return downloaded_files
+
+    def move_files(self, file_mappings: list[FileTransfer]) -> None:
+        assert self.bucket is not None, "call to move_files on uninitialized s3 session"
+
+        for src_file, dest_file in file_mappings:
+            if not isinstance(src_file, FileSystemFileReference) or not isinstance(
+                dest_file, FileSystemFileReference
+            ):
+                raise IncompatibleFileReference()
+            self.bucket.Object(dest_file.filepath).copy_from(
+                CopySource={
+                    "Bucket": self.bucket.name,
+                    "Key": src_file.filepath,
+                }
+            )
+            self.bucket.Object(src_file.filepath).delete()