UncountablePythonSDK 0.0.82__py3-none-any.whl → 0.0.132__py3-none-any.whl

pkgs/argument_parser/argument_parser.py

@@ -1,21 +1,25 @@
+from __future__ import annotations
+
 import dataclasses
+import datetime
 import math
 import types
 import typing
+from abc import ABC, abstractmethod
 from collections import defaultdict
-from dataclasses import MISSING, dataclass
-from datetime import date, datetime
+from datetime import date
 from decimal import Decimal
+from enum import Enum, auto
 from importlib import resources
 
 import dateutil.parser
+import msgspec.yaml
 
 from pkgs.serialization import (
     MissingSentryType,
     OpaqueKey,
     get_serial_class_data,
     get_serial_union_data,
-    yaml,
 )
 
 from ._is_enum import is_string_enum_class
@@ -27,13 +31,41 @@ ParserFunction = typing.Callable[[typing.Any], T]
 ParserCache = dict[type[typing.Any], ParserFunction[typing.Any]]
 
 
-@dataclass(frozen=True, eq=True)
+class SourceEncoding(Enum):
+    API = auto()
+    STORAGE = auto()
+
+
+@dataclasses.dataclass(frozen=True, eq=True)
 class ParserOptions:
-    convert_to_snake_case: bool
+    encoding: SourceEncoding
     strict_property_parsing: bool = False
 
+    @staticmethod
+    def Api(*, strict_property_parsing: bool = False) -> ParserOptions:
+        return ParserOptions(
+            encoding=SourceEncoding.API, strict_property_parsing=strict_property_parsing
+        )
+
+    @staticmethod
+    def Storage(*, strict_property_parsing: bool = False) -> ParserOptions:
+        return ParserOptions(
+            encoding=SourceEncoding.STORAGE,
+            strict_property_parsing=strict_property_parsing,
+        )
+
+    @property
+    def from_camel_case(self) -> bool:
+        return self.encoding == SourceEncoding.API
+
+    @property
+    def allow_direct_type(self) -> bool:
+        """This allows parsing from a DB column without having to check whether it's
+        the native format of the type, a JSON column, or a string encoding."""
+        return self.encoding == SourceEncoding.STORAGE
 
-@dataclass(frozen=True)
+
+@dataclasses.dataclass(frozen=True)
 class ParserContext:
     options: ParserOptions
     cache: ParserCache
@@ -52,20 +84,38 @@ class ParserExtraFieldsError(ParserError):
         return f"extra fields were provided: {', '.join(self.extra_fields)}"
 
 
+def is_union(field_type: typing.Any) -> bool:
+    origin = typing.get_origin(field_type)
+    return origin is typing.Union or origin is types.UnionType
+
+
 def is_optional(field_type: typing.Any) -> bool:
-    return typing.get_origin(field_type) is typing.Union and type(
-        None
-    ) in typing.get_args(field_type)
+    return is_union(field_type) and type(None) in typing.get_args(field_type)
 
 
 def is_missing(field_type: typing.Any) -> bool:
-    origin = typing.get_origin(field_type)
-    if origin is not typing.Union:
+    if not is_union(field_type):
         return False
     args = typing.get_args(field_type)
     return not (len(args) == 0 or args[0] is not MissingSentryType)
 
 
+def _has_field_default(field: dataclasses.Field[typing.Any]) -> bool:
+    return (
+        field.default != dataclasses.MISSING
+        and not isinstance(field.default, MissingSentryType)
+    ) or field.default_factory != dataclasses.MISSING
+
+
+def _get_field_default(
+    field: dataclasses.Field[typing.Any],
+) -> typing.Any:
+    if field.default != dataclasses.MISSING:
+        return field.default
+    assert field.default_factory != dataclasses.MISSING
+    return field.default_factory()
+
+
 def _invoke_tuple_parsers(
     tuple_type: type[T],
     arg_parsers: typing.Sequence[typing.Callable[[typing.Any], object]],
@@ -118,11 +168,39 @@ def _invoke_membership_parser(
     raise ValueError(f"Expected value from {expected_values} but got value {value}")
 
 
+# Uses `is` to compare
+def _build_identity_parser(
+    identity_value: T,
+) -> ParserFunction[T]:
+    def parse(value: typing.Any) -> T:
+        if value is identity_value:
+            return identity_value
+        raise ValueError(
+            f"Expected value {identity_value} (type: {type(identity_value)}) but got value {value} (type: {type(value)})"
+        )
+
+    return parse
+
+
+NONE_IDENTITY_PARSER = _build_identity_parser(None)
+
+
 def _build_parser_discriminated_union(
-    discriminator: str, discriminator_map: dict[str, ParserFunction[T]]
+    context: ParserContext,
+    discriminator_raw: str,
+    discriminator_map: dict[str, ParserFunction[T]],
 ) -> ParserFunction[T]:
+    discriminator = (
+        snake_to_camel_case(discriminator_raw)
+        if context.options.from_camel_case
+        else discriminator_raw
+    )
+
     def parse(value: typing.Any) -> typing.Any:
-        discriminant = value.get(discriminator)
+        if context.options.allow_direct_type and dataclasses.is_dataclass(value):
+            discriminant = getattr(value, discriminator)
+        else:
+            discriminant = value.get(discriminator)
         if discriminant is None:
             raise ValueError("missing-union-discriminant")
         if not isinstance(discriminant, str):
@@ -138,20 +216,10 @@ def _build_parser_discriminated_union(
 def _build_parser_inner(
     parsed_type: type[T],
     context: ParserContext,
-    *,
-    convert_string_to_snake_case: bool = False,
 ) -> ParserFunction[T]:
     """
-    convert_to_snake_case - internal flag
-                            if convert_to_snake_case is True, and parsed_type is str,
-                            then the generated parser will convert camel to snake case case
-                            should only be True for cases like dictionary keys
-                            should only be True if options.convert_to_snake_case is True
-
-    NOTE: This argument makes caching at this level difficult, as the cache-map
-    would need to vary based on this argument. For this reason only dataclasses
-    are cached now, as they don't use the argument, and they're known to be safe.
-    This is also enough to support some recursion.
+    IMPROVE: We can now cache at this level, to avoid producing redundant
+    internal parsers.
     """
 
     serial_union = get_serial_union_data(parsed_type)
@@ -163,6 +231,7 @@ def _build_parser_inner(
             parsed_type = serial_union.get_union_underlying()
         else:
             return _build_parser_discriminated_union(
+                context,
                 discriminator,
                 {
                     key: _build_parser_inner(value, context)
@@ -171,7 +240,7 @@ def _build_parser_inner(
             )
 
     if dataclasses.is_dataclass(parsed_type):
-        return _build_parser_dataclass(parsed_type, context)  # type: ignore[arg-type]
+        return _build_parser_dataclass(parsed_type, context)
 
     # namedtuple support
     if is_namedtuple_type(parsed_type):
@@ -184,15 +253,17 @@ def _build_parser_inner(
             field_name: field_parser(
                 value.get(
                     snake_to_camel_case(field_name)
-                    if context.options.convert_to_snake_case
+                    if context.options.from_camel_case
                     else field_name
                 )
             )
             for field_name, field_parser in field_parsers
         })
 
+    # IMPROVE: unclear why we need == here
     if parsed_type == type(None):  # noqa: E721
-        return lambda value: _invoke_membership_parser({None}, value)  # type: ignore
+        # Need to convince type checker that parsed_type is type(None)
+        return typing.cast(ParserFunction[T], NONE_IDENTITY_PARSER)
 
     origin = typing.get_origin(parsed_type)
     if origin is tuple:
@@ -219,7 +290,7 @@ def _build_parser_inner(
         arg_parsers = [_build_parser_inner(arg, context) for arg in sorted_args]
         return lambda value: _invoke_fallback_parsers(parsed_type, arg_parsers, value)
 
-    if parsed_type is typing.Any:  # type: ignore[comparison-overlap]
+    if parsed_type is typing.Any:
         return lambda value: value
 
     if origin in (list, set):
@@ -245,23 +316,32 @@ def _build_parser_inner(
         args = typing.get_args(parsed_type)
         if len(args) != 2:
             raise ValueError("Dict types only support two arguments for now")
-        k_parser = _build_parser_inner(
+        k_inner_parser = _build_parser_inner(
             args[0],
             context,
-            convert_string_to_snake_case=context.options.convert_to_snake_case,
         )
+
+        def key_parser(value: typing.Any) -> object:
+            inner = k_inner_parser(value)
+            if (
+                isinstance(inner, str)
+                # enum keys and OpaqueData's would also have string value types,
+                # but their explicit type is not a string, thus shouldn't be converted
+                and args[0] is str
+                and context.options.from_camel_case
+            ):
+                return camel_to_snake_case(value)
+            return inner
+
         v_parser = _build_parser_inner(args[1], context)
         return lambda value: origin(
-            (k_parser(k), v_parser(v)) for k, v in value.items()
+            (key_parser(k), v_parser(v)) for k, v in value.items()
         )
 
     if origin == typing.Literal:
         valid_values: set[T] = set(typing.get_args(parsed_type))
         return lambda value: _invoke_membership_parser(valid_values, value)
 
-    if parsed_type is str and convert_string_to_snake_case:
-        return lambda value: camel_to_snake_case(value)  # type: ignore
-
     if parsed_type is int:
         # first parse ints to decimal to allow scientific notation and decimals
         # e.g. (1) 1e4 => 1000, (2) 3.0 => 3
@@ -281,11 +361,25 @@ def _build_parser_inner(
 
         return parse_int
 
-    if parsed_type is datetime:
-        return lambda value: dateutil.parser.isoparse(value)  # type:ignore
+    if parsed_type is datetime.datetime:
+
+        def parse_datetime(value: typing.Any) -> T:
+            if context.options.allow_direct_type and isinstance(
+                value, datetime.datetime
+            ):
+                return value  # type: ignore
+            return dateutil.parser.isoparse(value)  # type:ignore
+
+        return parse_datetime
 
     if parsed_type is date:
-        return lambda value: date.fromisoformat(value)  # type:ignore
+
+        def parse_date(value: typing.Any) -> T:
+            if context.options.allow_direct_type and isinstance(value, date):
+                return value  # type:ignore
+            return date.fromisoformat(value)  # type:ignore
+
+        return parse_date
 
     # MyPy: It's unclear why `parsed_type in (str, OpaqueKey)` is flagged as invalid
     # Thus an or statement is used instead, which isn't flagged as invalid.
@@ -320,7 +414,17 @@ def _build_parser_inner(
             raise ValueError("Missing type cannot be parsed directly")
 
         return error
-    raise ValueError(f"Unhandled type {parsed_type}")
+
+    # Check last for generic annotated types and process them unwrapped
+    # this must be last, since some of the expected types, like Unions,
+    # will also be annotated, but have a special form
+    if typing.get_origin(parsed_type) is typing.Annotated:
+        return _build_parser_inner(
+            parsed_type.__origin__,  # type: ignore[attr-defined]
+            context,
+        )
+
+    raise ValueError(f"Unhandled type {parsed_type}/{origin}")
 
 
 def _build_parser_dataclass(
@@ -335,8 +439,7 @@ def _build_parser_dataclass(
     cur_parser = context.cache.get(parsed_type)
     if cur_parser is not None:
         return cur_parser
-
-    type_hints = typing.get_type_hints(parsed_type)
+    type_hints = typing.get_type_hints(parsed_type, include_extras=True)
     dc_field_parsers: list[
         tuple[
             dataclasses.Field[typing.Any],
@@ -351,29 +454,31 @@ def _build_parser_dataclass(
         return (
             snake_to_camel_case(field_name)
             if (
-                context.options.convert_to_snake_case
+                context.options.from_camel_case
                 and not serial_class_data.has_unconverted_key(field_name)
             )
             else field_name
         )
 
     def parse(value: typing.Any) -> typing.Any:
+        # Use an exact type match to prevent base/derived class mismatches
+        if context.options.allow_direct_type and type(value) is parsed_type:
+            return value
+
         data: dict[typing.Any, typing.Any] = {}
         for field, field_type, field_parser in dc_field_parsers:
             field_raw_value = None
             try:
                 field_raw_value = value.get(
                     resolve_serialized_field_name(field_name=field.name),
-                    MISSING,
+                    dataclasses.MISSING,
                 )
                 field_value: typing.Any
-                if field_raw_value == MISSING:
+                if field_raw_value == dataclasses.MISSING:
                     if serial_class_data.has_parse_require(field.name):
                         raise ValueError("missing-required-field", field.name)
-                    if field.default != MISSING:
-                        field_value = field.default
-                    elif field.default_factory != MISSING:
-                        field_value = field.default_factory()
+                    if _has_field_default(field):
+                        field_value = _get_field_default(field)
                     elif is_missing(field_type):
                         field_value = MissingSentryType()
                     elif is_optional(field_type):
@@ -385,6 +490,13 @@ def _build_parser_dataclass(
                         field_value = False
                     else:
                         raise ValueError("missing-value-for-field", field.name)
+                elif (
+                    field_raw_value is None
+                    and not is_optional(field_type)
+                    and _has_field_default(field)
+                    and not serial_class_data.has_parse_require(field.name)
+                ):
+                    field_value = _get_field_default(field)
                 elif serial_class_data.has_unconverted_value(field.name):
                     field_value = field_raw_value
                 else:
@@ -452,15 +564,46 @@ def build_parser(
     return built_parser
 
 
-class CachedParser(typing.Generic[T]):
+class ParserBase(ABC, typing.Generic[T]):
+    def parse_from_encoding(
+        self,
+        args: typing.Any,
+        *,
+        source_encoding: SourceEncoding,
+    ) -> T:
+        match source_encoding:
+            case SourceEncoding.API:
+                return self.parse_api(args)
+            case SourceEncoding.STORAGE:
+                return self.parse_storage(args)
+            case _:
+                typing.assert_never(source_encoding)
+
+    # IMPROVE: Args would be better typed as "object"
+    @abstractmethod
+    def parse_storage(self, args: typing.Any) -> T: ...
+
+    @abstractmethod
+    def parse_api(self, args: typing.Any) -> T: ...
+
+    def parse_yaml_file(self, path: str) -> T:
+        with open(path, encoding="utf-8") as data_in:
+            return self.parse_storage(msgspec.yaml.decode(data_in.read()))
+
+    def parse_yaml_resource(self, package: resources.Package, resource: str) -> T:
+        with resources.open_text(package, resource) as fp:
+            return self.parse_storage(msgspec.yaml.decode(fp.read()))
+
+
+class CachedParser(ParserBase[T], typing.Generic[T]):
     def __init__(
         self,
         args: type[T],
         strict_property_parsing: bool = False,
     ):
         self.arguments = args
-        self.parser_api: typing.Optional[ParserFunction[T]] = None
-        self.parser_storage: typing.Optional[ParserFunction[T]] = None
+        self.parser_api: ParserFunction[T] | None = None
+        self.parser_storage: ParserFunction[T] | None = None
         self.strict_property_parsing = strict_property_parsing
 
     def parse_api(self, args: typing.Any) -> T:
@@ -474,8 +617,7 @@ class CachedParser(typing.Generic[T]):
         if self.parser_api is None:
             self.parser_api = build_parser(
                 self.arguments,
-                ParserOptions(
-                    convert_to_snake_case=True,
+                ParserOptions.Api(
                     strict_property_parsing=self.strict_property_parsing,
                 ),
             )
@@ -489,18 +631,9 @@ class CachedParser(typing.Generic[T]):
         if self.parser_storage is None:
             self.parser_storage = build_parser(
                 self.arguments,
-                ParserOptions(
-                    convert_to_snake_case=False,
+                ParserOptions.Storage(
                     strict_property_parsing=self.strict_property_parsing,
                 ),
             )
         assert self.parser_storage is not None
         return self.parser_storage(args)
-
-    def parse_yaml_file(self, path: str) -> T:
-        with open(path, encoding="utf-8") as data_in:
-            return self.parse_storage(yaml.safe_load(data_in))
-
-    def parse_yaml_resource(self, package: resources.Package, resource: str) -> T:
-        with resources.open_text(package, resource) as fp:
-            return self.parse_storage(yaml.safe_load(fp))

pkgs/filesystem_utils/__init__.py

@@ -17,3 +17,4 @@ 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

@@ -1,6 +1,6 @@
 import os
 from io import BytesIO
-from typing import Any, Optional
+from typing import Any
 
 from google.oauth2 import service_account
 from googleapiclient.discovery import build as build_gdrive_connection
@@ -30,7 +30,7 @@ def download_gdrive_file(
     mime_type: str,
     *,
     verbose: bool = False,
-) -> Optional[FileObjectData]:
+) -> FileObjectData | None:
     if "folder" in mime_type:
         if verbose:
             print(f"{filename} is a folder and will not be downloaded.")
@@ -63,7 +63,7 @@ def download_gdrive_file(
     downloader = MediaIoBaseDownload(file_handler, file_request)
     download_complete = False
     while not download_complete:
-        status, download_complete = downloader.next_chunk()
+        _status, download_complete = downloader.next_chunk()
 
     file_handler.seek(0)
     file_data = file_handler.read()
@@ -148,7 +148,7 @@ def move_gdrive_file(
     src_file_id: str,
     dest_folder_id: str,
     *,
-    dest_filename: Optional[str] = None,
+    dest_filename: str | None = None,
 ) -> None:
     # Retrieve the existing parents to remove
     file = (
@@ -197,7 +197,7 @@ class GDriveSession(FileSystemSession):
         dir_path: FileSystemObject,
         *,
         recursive: bool = False,
-        valid_file_extensions: Optional[tuple[str, ...]] = None,
+        valid_file_extensions: tuple[str, ...] | None = None,
     ) -> list[FileSystemObject]:
         if not isinstance(dir_path, RemoteObjectReference):
             raise IncompatibleFileReference(

pkgs/filesystem_utils/_s3_session.py

@@ -1,6 +1,7 @@
 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,
@@ -37,7 +38,7 @@ class S3Session(FileSystemSession):
             aws_session_token=self.config.session_token,
         )
 
-        self.bucket = s3_resource.Bucket(self.config.bucket_name)
+        self.bucket: Bucket | None = s3_resource.Bucket(self.config.bucket_name)
 
     def __enter__(self) -> "S3Session":
         self.start()

pkgs/filesystem_utils/_sftp_session.py

@@ -1,7 +1,6 @@
 import os
 from collections.abc import Iterable
 from io import BytesIO
-from typing import Optional
 
 import paramiko
 import pysftp
@@ -30,8 +29,8 @@ def list_sftp_files(
     connection: pysftp.Connection,
     dir_path: str,
     *,
-    valid_extensions: Optional[Iterable[str]] = None,
-    parent_dir_path: Optional[str] = None,
+    valid_extensions: Iterable[str] | None = None,
+    parent_dir_path: str | None = None,
     recursive: bool = True,
 ) -> list[str]:
     file_paths: list[str] = []
@@ -55,6 +54,10 @@ def list_sftp_files(
             os.path.join(dir_path, file)
             for file in connection.listdir(dir_path)
             if connection.isfile(os.path.join(dir_path, file))
+            and (
+                valid_extensions is None
+                or os.path.splitext(file)[1] in valid_extensions
+            )
         ])
     return file_paths