UncountablePythonSDK 0.0.8__py3-none-any.whl → 0.0.92__py3-none-any.whl

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,117 @@
+from io import BytesIO
+
+from boto3.session import Session
+
+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 = 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()

pkgs/filesystem_utils/_sftp_session.py

@@ -0,0 +1,147 @@
+import os
+from collections.abc import Iterable
+from io import BytesIO
+
+import paramiko
+import pysftp
+
+from pkgs.filesystem_utils.file_type_utils import (
+    FileObjectData,
+    FileSystemFileReference,
+    FileSystemObject,
+    FileSystemSFTPConfig,
+    FileTransfer,
+    IncompatibleFileReference,
+)
+
+from .filesystem_session import FileSystemSession
+
+
+def move_sftp_files(
+    connection: pysftp.Connection,
+    src_filepath: str,
+    dest_filepath: str,
+) -> None:
+    connection.rename(src_filepath, dest_filepath)
+
+
+def list_sftp_files(
+    connection: pysftp.Connection,
+    dir_path: str,
+    *,
+    valid_extensions: Iterable[str] | None = None,
+    parent_dir_path: str | None = None,
+    recursive: bool = True,
+) -> list[str]:
+    file_paths: list[str] = []
+    if recursive:
+
+        def _skip(name: str) -> None:
+            return
+
+        def _add_file(path: str) -> None:
+            if (
+                valid_extensions is None
+                or os.path.splitext(path)[1] in valid_extensions
+            ) and (parent_dir_path is None or os.path.dirname(path) == parent_dir_path):
+                file_paths.append(path)
+
+        connection.walktree(
+            dir_path, fcallback=_add_file, dcallback=_skip, ucallback=_skip
+        )
+    else:
+        file_paths.extend([
+            os.path.join(dir_path, file)
+            for file in connection.listdir(dir_path)
+            if connection.isfile(os.path.join(dir_path, file))
+        ])
+    return file_paths
+
+
+class SFTPSession(FileSystemSession):
+    def __init__(self, sftp_config: FileSystemSFTPConfig) -> None:
+        super().__init__()
+        self.host: str = sftp_config.ip
+        self.username: str = sftp_config.username
+        self.key_file: str | paramiko.RSAKey | None = (
+            sftp_config.pem_path
+            if sftp_config.pem_path is not None
+            else sftp_config.pem_key
+        )
+        self.password: str | None = sftp_config.password
+
+    def start(self) -> None:
+        cnopts = pysftp.CnOpts()
+        cnopts.hostkeys = None
+        if self.key_file is not None:
+            self.connection = pysftp.Connection(
+                self.host,
+                username=self.username,
+                private_key=self.key_file,
+                cnopts=cnopts,
+            )
+        elif self.password is not None:
+            self.connection = pysftp.Connection(
+                self.host,
+                username=self.username,
+                password=self.password,
+                cnopts=cnopts,
+            )
+        else:
+            raise pysftp.CredentialException(
+                "Must specify either a private key path or a password."
+            )
+
+    def __enter__(self) -> "SFTPSession":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        self.connection.close()
+
+    def list_files(
+        self,
+        dir_path: FileSystemObject,
+        *,
+        recursive: bool = True,
+        valid_extensions: list[str] | None = None,
+    ) -> list[FileSystemObject]:
+        if not isinstance(
+            dir_path, FileSystemFileReference
+        ) or not self.connection.isdir(dir_path.filepath):
+            raise IncompatibleFileReference()
+
+        return [
+            FileSystemFileReference(file_path)
+            for file_path in list_sftp_files(
+                self.connection,
+                dir_path.filepath,
+                recursive=recursive,
+                valid_extensions=valid_extensions,
+            )
+        ]
+
+    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()
+            filepath = file_object.filepath
+            file_data = self.connection.open(filepath).read()
+            downloaded_file = FileObjectData(
+                file_data, BytesIO(file_data), file_object.filename, filepath=filepath
+            )
+            if downloaded_file is not None:
+                downloaded_files.append(downloaded_file)
+        return downloaded_files
+
+    def move_files(self, file_mappings: list[FileTransfer]) -> None:
+        for src_file, dest_file in file_mappings:
+            if not isinstance(src_file, FileSystemFileReference) or not isinstance(
+                dest_file, FileSystemFileReference
+            ):
+                raise IncompatibleFileReference()
+            move_sftp_files(self.connection, src_file.filepath, dest_file.filepath)