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

pkgs/type_spec/value_spec/__main__.py

@@ -13,11 +13,14 @@ One of the following can be specified on the name of a argument:
 After that you can also specify a `!` indicating the argument may not be null.
 If this is not specified, then a null input on this argument should produce a null output.
 We prefer not to use `!` as we want to encourage null pass-through where possible.
-If null is allowed as a legitimate value, such as in conditionals like `if`, then `!` must be specified.
+
+If null is allowed as a legitimate value, such as in conditionals like `is_null`,
+then `!usenull` must be specified, this distinguishes it from the pass-through case.
+The accepted argument type must accept "None", it is not implied.
 """
 
 import sys
-from typing import TypeVar, cast
+from typing import Match, Pattern, TypeVar, cast
 
 import regex as re
 
@@ -53,7 +56,7 @@ class Source:
     def has_more(self) -> bool:
         return self._at < len(self._text)
 
-    def match(self, expression: re.Pattern) -> re.Match | None:
+    def match(self, expression: Pattern[str]) -> Match[str] | None:
         self.skip_space()
         m = expression.match(self._text, self._at)
         if m is not None:
@@ -84,7 +87,7 @@ class Source:
         return self._text[start : self._at]
 
 
-_re_argument_name = re.compile(r"([a-z_]+)(\?|\+)?(!)?:")
+_re_argument_name = re.compile(r"([a-z_]+)(\?|\+)?(!|!usenull)?:")
 
 
 def parse_function_signature(text: str) -> ParsedFunctionSignature:
@@ -101,11 +104,18 @@ def parse_function_signature(text: str) -> ParsedFunctionSignature:
 
         type_str = source.extract_type()
         ref_name = arg_group.group(1)
-        is_missing = arg_group.group(2) == "?"
-        is_repeating = arg_group.group(2) == "+"
-        pass_null = arg_group.group(3) is None
+        # is_missing = arg_group.group(2) == "?"
+        # is_repeating = arg_group.group(2) == "+"
         type_path = parse_type_str(type_str)
 
+        match arg_group.group(3):
+            case "!":
+                on_null = value_spec_t.OnNull.DISALLOW
+            case "!usenull":
+                on_null = value_spec_t.OnNull.USE
+            case _:
+                on_null = value_spec_t.OnNull.PASS
+
         extant = value_spec_t.ArgumentExtant.REQUIRED
         extant_marker = arg_group.group(2)
         if extant_marker == "?":
@@ -116,7 +126,7 @@ def parse_function_signature(text: str) -> ParsedFunctionSignature:
         arguments.append(
             ParsedFunctionArgument(
                 ref_name=ref_name,
-                pass_null=pass_null,
+                on_null=on_null,
                 extant=extant,
                 type_path=type_path,
             )
@@ -145,6 +155,7 @@ key_return = "return"
 key_description = "description"
 key_brief = "brief"
 key_name = "name"
+key_draft = "draft"
 
 
 TypeT = TypeVar("TypeT")
@@ -208,7 +219,7 @@ def main() -> None:
                         name=arg_name,
                         description=arg_description,
                         type=convert_to_value_spec_type(in_argument.type_path),
-                        pass_null=in_argument.pass_null,
+                        on_null=in_argument.on_null,
                         extant=in_argument.extant,
                     )
                 )
@@ -219,6 +230,11 @@ def main() -> None:
 
             brief = get_as(spec, key_brief, str)
             description = get_as(spec, key_description, str)
+            draft = (
+                get_as(spec, key_draft, bool)
+                if spec.get(key_draft) is not None
+                else None
+            )
 
             return_value = get(spec, key_return)
             where.append("return")
@@ -235,6 +251,7 @@ def main() -> None:
                         type=convert_to_value_spec_type(parsed.return_type_path),
                         description=return_description,
                     ),
+                    draft=draft,
                 )
             )
             where.pop()

pkgs/type_spec/value_spec/convert_type.py

@@ -25,10 +25,17 @@ TYPE_MAP = {
     "List": MappedType(base_type=value_spec_t.BaseType.LIST, param_count=1),
     "Optional": MappedType(base_type=value_spec_t.BaseType.OPTIONAL, param_count=1),
     "String": MappedType(base_type=value_spec_t.BaseType.STRING),
-    "Union": MappedType(base_type=value_spec_t.BaseType.UNION, variable_param_count=True),
+    "Union": MappedType(
+        base_type=value_spec_t.BaseType.UNION, variable_param_count=True
+    ),
     # not part of type_spec's types now
     "Symbol": MappedType(base_type=value_spec_t.BaseType.SYMBOL),
     "Any": MappedType(base_type=value_spec_t.BaseType.ANY),
+    "None": MappedType(base_type=value_spec_t.BaseType.NONE),
+    "Tuple": MappedType(
+        base_type=value_spec_t.BaseType.TUPLE, variable_param_count=True
+    ),
+    "Never": MappedType(base_type=value_spec_t.BaseType.NEVER),
 }
 
 
@@ -54,3 +61,16 @@ def convert_to_value_spec_type(parsed: ParsedTypePath) -> value_spec_t.ValueType
     return value_spec_t.ValueType(base_type=mapped.base_type, parameters=parameters)
 
     # Our formatter was duplicating the previous line for an unknown reason, this comment blocks that
+
+
+def convert_from_value_spec_type(
+    base_type: value_spec_t.BaseType,
+) -> str:
+    for type_spec_type, mapped_type in TYPE_MAP.items():
+        if (
+            mapped_type.base_type == base_type
+            and mapped_type.param_count == 0
+            and mapped_type.variable_param_count is False
+        ):
+            return type_spec_type
+    raise ValueError(f"invalid value spec type {base_type}")

pkgs/type_spec/value_spec/emit_python.py

@@ -70,14 +70,19 @@ def _emit_function_wrapper(function: value_spec_t.Function) -> str:
         else:
             python_type = _emit_python_type(argument.type)
             if (
-                argument.pass_null
+                argument.on_null == value_spec_t.OnNull.PASS
                 or argument.extant == value_spec_t.ArgumentExtant.MISSING
             ):
                 python_type += " | None"
                 any_pass_null = True
+
+            if python_type.startswith("base_t.ExtJsonValue"):
+                return_statement = f"self._extract({index})"
+            else:
+                return_statement = f"cast({python_type}, self._extract({index}))"
             out.write(
                 f"""{INDENT}def get_{argument.ref_name}(self) -> {python_type}:
-{INDENT}{INDENT}return cast({python_type}, self._extract({index}))
+{INDENT}{INDENT}return {return_statement}
 """
             )
         out.write("\n")
@@ -160,8 +165,12 @@ def _emit_function(function: value_spec_t.Function, indent: str) -> str:
     sub_indent = indent + INDENT
     out.write(f"{_function_symbol_name(function)} = value_spec_t.Function(\n")
     out.write(f"{sub_indent}name={encode_common_string(function.name)},\n")
-    out.write(f"{sub_indent}description={encode_common_string(function.description)},\n")
+    out.write(
+        f"{sub_indent}description={encode_common_string(function.description)},\n"
+    )
     out.write(f"{sub_indent}brief={encode_common_string(function.brief)},\n")
+    if function.draft:
+        out.write(f"{sub_indent}draft={function.draft},\n")
     out.write(
         f"{sub_indent}return_value={_emit_function_return(function.return_value, sub_indent)},\n"
     )
@@ -184,16 +193,25 @@ def _emit_argument(argument: value_spec_t.FunctionArgument, indent: str) -> str:
     out.write("value_spec_t.FunctionArgument(\n")
     out.write(f"{sub_indent}ref_name={encode_common_string(argument.ref_name)},\n")
     out.write(f"{sub_indent}name={encode_common_string(argument.name)},\n")
-    out.write(f"{sub_indent}description={encode_common_string(argument.description)},\n")
-    out.write(f"{sub_indent}pass_null={str(argument.pass_null)},\n")
-    out.write(f"{sub_indent}extant=value_spec_t.ArgumentExtant.{argument.extant.name},\n")
+    out.write(
+        f"{sub_indent}description={encode_common_string(argument.description)},\n"
+    )
+    # Quick enum emit since we have only one such type here
+    out.write(
+        f"{sub_indent}on_null=value_spec_t.OnNull.{str(argument.on_null).upper()},\n"
+    )
+    out.write(
+        f"{sub_indent}extant=value_spec_t.ArgumentExtant.{argument.extant.name},\n"
+    )
     out.write(f"{sub_indent}type={_emit_type(argument.type, sub_indent)},\n")
     out.write(f"{indent})")
 
     return out.getvalue()
 
 
-def _emit_function_return(return_value: value_spec_t.FunctionReturn, indent: str) -> str:
+def _emit_function_return(
+    return_value: value_spec_t.FunctionReturn, indent: str
+) -> str:
     out = io.StringIO()
 
     sub_indent = indent + INDENT

pkgs/type_spec/value_spec/types.py

@@ -8,7 +8,7 @@ from ..util import ParsedTypePath
 @dataclass(kw_only=True, frozen=True)
 class ParsedFunctionArgument:
     ref_name: str
-    pass_null: bool
+    on_null: value_spec_t.OnNull
     extant: value_spec_t.ArgumentExtant
     type_path: ParsedTypePath
 

uncountable/core/async_batch.py

@@ -4,7 +4,7 @@ from uncountable.types.async_batch import AsyncBatchRequest
 from uncountable.types.async_batch_processor import AsyncBatchProcessorBase
 
 
-class AsyncBatchSubmissionError(BaseException):
+class AsyncBatchSubmissionError(Exception):
     pass
 
 

uncountable/core/client.py

@@ -1,21 +1,24 @@
 import base64
-import json
+import datetime
+import re
 import typing
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import UTC, timedelta
 from enum import StrEnum
-from urllib.parse import urljoin
+from io import BytesIO
+from urllib.parse import unquote, urljoin
 from uuid import uuid4
 
 import requests
+import simplejson as json
 from opentelemetry.sdk.resources import Attributes
 from requests.exceptions import JSONDecodeError
 
 from pkgs.argument_parser import CachedParser
-from pkgs.serialization_util import serialize_for_api
-from pkgs.serialization_util.serialization_helpers import JsonValue
-from uncountable.core.version import get_version
-from uncountable.integration.telemetry import JobLogger
+from pkgs.serialization_util import JsonValue, serialize_for_api
+from uncountable.core.environment import get_version
+from uncountable.integration.telemetry import Logger, push_scope_optional
+from uncountable.types import download_file_t
 from uncountable.types.client_base import APIRequest, ClientMethods
 from uncountable.types.client_config import ClientConfigOptions
 
@@ -41,14 +44,14 @@ class HTTPRequestBase:
 
 @dataclass(kw_only=True)
 class HTTPGetRequest(HTTPRequestBase):
-    method = EndpointMethod.GET
+    method: EndpointMethod = EndpointMethod.GET
     query_params: dict[str, str]
 
 
 @dataclass(kw_only=True)
 class HTTPPostRequest(HTTPRequestBase):
-    method = EndpointMethod.POST
-    body: typing.Union[str, dict[str, str]]
+    method: EndpointMethod = EndpointMethod.POST
+    body: str | dict[str, str]
 
 
 HTTPRequest = HTTPPostRequest | HTTPGetRequest
@@ -56,29 +59,39 @@ HTTPRequest = HTTPPostRequest | HTTPGetRequest
 
 @dataclass(kw_only=True)
 class ClientConfig(ClientConfigOptions):
-    transform_request: typing.Callable[[requests.Request], requests.Request] | None = None
-    job_logger: typing.Optional[JobLogger] = None
+    transform_request: typing.Callable[[requests.Request], requests.Request] | None = (
+        None
+    )
+    logger: Logger | None = None
 
 
 OAUTH_REFRESH_WINDOW_SECONDS = 60 * 5
 
 
-class APIResponseError(BaseException):
+class APIResponseError(Exception):
     status_code: int
     message: str
     extra_details: dict[str, JsonValue] | None
 
     def __init__(
-        self, status_code: int, message: str, extra_details: dict[str, JsonValue] | None
+        self,
+        status_code: int,
+        message: str,
+        extra_details: dict[str, JsonValue] | None,
+        request_id: str,
     ) -> None:
         super().__init__(status_code, message, extra_details)
         self.status_code = status_code
         self.message = message
         self.extra_details = extra_details
+        self.request_id = request_id
 
     @classmethod
     def construct_error(
-        cls, status_code: int, extra_details: dict[str, JsonValue] | None
+        cls,
+        status_code: int,
+        extra_details: dict[str, JsonValue] | None,
+        request_id: str,
     ) -> "APIResponseError":
         message: str
         match status_code:
@@ -101,11 +114,23 @@ class APIResponseError(BaseException):
             case _:
                 message = "unknown error"
         return APIResponseError(
-            status_code=status_code, message=message, extra_details=extra_details
+            status_code=status_code,
+            message=message,
+            extra_details=extra_details,
+            request_id=request_id,
         )
 
+    def __str__(self) -> str:
+        details_obj = {
+            "request_id": self.request_id,
+            "status_code": self.status_code,
+            "extra_details": self.extra_details,
+        }
+        details = json.dumps(details_obj)
+        return f"API response error ({self.status_code}): '{self.message}'. Details: {details}"
 
-class SDKError(BaseException):
+
+class SDKError(Exception):
     message: str
     request_id: str
 
@@ -121,7 +146,7 @@ class SDKError(BaseException):
 @dataclass(kw_only=True)
 class OAuthBearerTokenCache:
     token: str
-    expires_at: datetime
+    expires_at: datetime.datetime
 
 
 @dataclass(kw_only=True)
@@ -135,6 +160,16 @@ class GetOauthBearerTokenData:
 oauth_bearer_token_data_parser = CachedParser(GetOauthBearerTokenData)
 
 
+@dataclass
+class DownloadedFile:
+    name: str
+    size: int
+    data: BytesIO
+
+
+DownloadedFiles = list[DownloadedFile]
+
+
 class Client(ClientMethods):
     _parser_map: dict[type, CachedParser] = {}
     _auth_details: AuthDetailsAll
@@ -153,37 +188,53 @@ class Client(ClientMethods):
     ):
         self._auth_details = auth_details
         self._base_url = base_url
-        self._file_uploader = FileUploader(self._base_url, self._auth_details)
         self._cfg = config or ClientConfig()
         self._session = requests.Session()
         self._session.verify = not self._cfg.allow_insecure_tls
+        self._file_uploader = FileUploader(
+            self._base_url,
+            self._auth_details,
+            self._cfg.allow_insecure_tls,
+            logger=self._cfg.logger,
+        )
 
-    def _get_response_json(
-        self, response: requests.Response, request_id: str
-    ) -> dict[str, JsonValue]:
+    @classmethod
+    def _validate_response_status(
+        cls, response: requests.Response, request_id: str
+    ) -> None:
         if response.status_code < 200 or response.status_code > 299:
             extra_details: dict[str, JsonValue] | None = None
             try:
                 data = response.json()
-                if "error" in data:
-                    extra_details = data
+                extra_details = data
             except JSONDecodeError:
-                pass
+                extra_details = {
+                    "body": response.text,
+                }
             raise APIResponseError.construct_error(
-                status_code=response.status_code, extra_details=extra_details
+                status_code=response.status_code,
+                extra_details=extra_details,
+                request_id=request_id,
             )
+
+    def _get_response_json(
+        self, response: requests.Response, request_id: str
+    ) -> dict[str, JsonValue]:
+        self._validate_response_status(response, request_id)
         try:
             return typing.cast(dict[str, JsonValue], response.json())
         except JSONDecodeError as e:
             raise SDKError("unable to process response", request_id=request_id) from e
 
-    def _send_request(self, request: requests.Request) -> requests.Response:
+    def _send_request(
+        self, request: requests.Request, *, timeout: float | None = None
+    ) -> requests.Response:
         if self._cfg.extra_headers is not None:
             request.headers = {**request.headers, **self._cfg.extra_headers}
         if self._cfg.transform_request is not None:
             request = self._cfg.transform_request(request)
         prepared_request = request.prepare()
-        response = self._session.send(prepared_request)
+        response = self._session.send(prepared_request, timeout=timeout)
         return response
 
     def do_request(self, *, api_request: APIRequest, return_type: type[DT]) -> DT:
@@ -201,15 +252,19 @@ class Client(ClientMethods):
             case _:
                 typing.assert_never(http_request)
         request.headers = http_request.headers
-        if self._cfg.job_logger is not None:
-            attributes: Attributes = {
-                "method": http_request.method,
-                "endpoint": api_request.endpoint,
-            }
-            with self._cfg.job_logger.push_scope("api_call", attributes=attributes):
-                response = self._send_request(request)
-        else:
-            response = self._send_request(request)
+        attributes: Attributes = {
+            "method": http_request.method,
+            "endpoint": api_request.endpoint,
+        }
+        with push_scope_optional(self._cfg.logger, "api_call", attributes=attributes):
+            if self._cfg.logger is not None:
+                self._cfg.logger.log_info(api_request.endpoint, attributes=attributes)
+            timeout = (
+                api_request.request_options.timeout_secs
+                if api_request.request_options is not None
+                else None
+            )
+            response = self._send_request(request, timeout=timeout)
         response_data = self._get_response_json(response, request_id=request_id)
         cached_parser = self._get_cached_parser(return_type)
         try:
@@ -227,7 +282,8 @@ class Client(ClientMethods):
         if (
             self._oauth_bearer_token_cache is None
             or (
-                self._oauth_bearer_token_cache.expires_at - datetime.now()
+                self._oauth_bearer_token_cache.expires_at
+                - datetime.datetime.now(tz=UTC)
             ).total_seconds()
             < OAUTH_REFRESH_WINDOW_SECONDS
         ):
@@ -243,7 +299,8 @@ class Client(ClientMethods):
             token_data = oauth_bearer_token_data_parser.parse_storage(data)
             self._oauth_bearer_token_cache = OAuthBearerTokenCache(
                 token=token_data.access_token,
-                expires_at=datetime.now() + timedelta(seconds=token_data.expires_in),
+                expires_at=datetime.datetime.now(tz=UTC)
+                + timedelta(seconds=token_data.expires_in),
             )
 
         return self._oauth_bearer_token_cache.token
@@ -286,6 +343,52 @@ class Client(ClientMethods):
             case _:
                 raise ValueError(f"unsupported request method: {method}")
 
+    def _get_downloaded_filename(self, *, cd: str | None) -> str:
+        if not cd:
+            return "Unknown"
+
+        fname = re.findall(r"filename\*=UTF-8''(.+)", cd)
+        if fname:
+            return unquote(fname[0])
+
+        fname = re.findall(r'filename="?(.+)"?', cd)
+        if fname:
+            return str(fname[0].strip('"'))
+
+        return "Unknown"
+
+    def download_files(
+        self, *, file_query: download_file_t.FileDownloadQuery
+    ) -> DownloadedFiles:
+        """Download a file from uncountable."""
+        request_id = str(uuid4())
+        api_request = APIRequest(
+            method=download_file_t.ENDPOINT_METHOD,
+            endpoint=download_file_t.ENDPOINT_PATH,
+            args=download_file_t.Arguments(
+                file_query=file_query,
+            ),
+        )
+        http_request = self._build_http_request(
+            api_request=api_request, request_id=request_id
+        )
+        request = requests.Request(http_request.method.value, http_request.url)
+        request.headers = http_request.headers
+        assert isinstance(http_request, HTTPGetRequest)
+        request.params = http_request.query_params
+        response = self._send_request(request)
+        self._validate_response_status(response, request_id)
+
+        content = response.content
+        content_disposition = response.headers.get("Content-Disposition", None)
+        return [
+            DownloadedFile(
+                name=self._get_downloaded_filename(cd=content_disposition),
+                size=len(content),
+                data=BytesIO(content),
+            )
+        ]
+
     def upload_files(
         self: typing.Self, *, file_uploads: list[FileUpload]
     ) -> list[UploadedFile]:

uncountable/core/environment.py

@@ -0,0 +1,41 @@
+import functools
+import os
+from importlib.metadata import PackageNotFoundError, version
+
+from uncountable.types import integration_server_t
+
+
+@functools.cache
+def get_version() -> str:
+    try:
+        version_str = version("UncountablePythonSDK")
+    except PackageNotFoundError:
+        version_str = "unknown"
+    return version_str
+
+
+def get_server_env() -> str | None:
+    return os.environ.get("UNC_SERVER_ENV")
+
+
+def get_http_server_port() -> int:
+    return int(os.environ.get("UNC_WEBHOOK_SERVER_PORT", "5001"))
+
+
+def get_local_admin_server_port() -> int:
+    return int(os.environ.get("UNC_ADMIN_SERVER_PORT", "50051"))
+
+
+def get_otel_enabled() -> bool:
+    return os.environ.get("UNC_OTEL_ENABLED") == "true"
+
+
+def get_profiles_module() -> str:
+    return os.environ["UNC_PROFILES_MODULE"]
+
+
+def get_integration_envs() -> list[integration_server_t.IntegrationEnvironment]:
+    return [
+        integration_server_t.IntegrationEnvironment(env)
+        for env in os.environ.get("UNC_INTEGRATION_ENVS", "prod").split(",")
+    ]

uncountable/core/file_upload.py

@@ -4,11 +4,13 @@ from dataclasses import dataclass
 from enum import StrEnum
 from io import BytesIO
 from pathlib import Path
-from typing import Generator, Literal, Self
+from typing import Generator, Literal, Self, assert_never
 
 import aiohttp
 import aiotus
 
+from uncountable.integration.telemetry import Logger, push_scope_optional
+
 from .types import AuthDetailsAll, AuthDetailsApiKey
 
 _CHUNK_SIZE = 5 * 1024 * 1024  # s3 requires 5MiB minimum
@@ -68,10 +70,19 @@ class UploadFailed(Exception):
 class FileUploader:
     _auth_details: AuthDetailsAll
     _base_url: str
-
-    def __init__(self: Self, base_url: str, auth_details: AuthDetailsAll) -> None:
+    _allow_insecure_tls: bool
+
+    def __init__(
+        self: Self,
+        base_url: str,
+        auth_details: AuthDetailsAll,
+        allow_insecure_tls: bool = False,
+        logger: Logger | None = None,
+    ) -> None:
         self._base_url = base_url
         self._auth_details = auth_details
+        self._allow_insecure_tls = allow_insecure_tls
+        self._logger = logger
 
     async def _upload_file(self: Self, file_upload: FileUpload) -> UploadedFile:
         creation_url = f"{self._base_url}/api/external/file_upload/files"
@@ -86,21 +97,44 @@ class FileUploader:
                 auth=auth, headers={"Origin": self._base_url}
             ) as session,
         ):
-            with file_upload_data(file_upload) as file_bytes:
-                location = await aiotus.upload(
-                    creation_url,
-                    file_bytes.bytes_data,
-                    {"filename": file_bytes.name.encode()},
-                    client_session=session,
-                    chunksize=_CHUNK_SIZE,
-                )
-                if location is None:
-                    raise UploadFailed(f"Failed to upload: {file_bytes.name}")
-                return UploadedFile(
-                    name=file_bytes.name, file_id=int(location.path.split("/")[-1])
-                )
-
-    def upload_files(self: Self, *, file_uploads: list[FileUpload]) -> list[UploadedFile]:
+            attributes = {}
+            match file_upload:
+                case MediaFileUpload():
+                    attributes["file_path"] = file_upload.path
+                case DataFileUpload():
+                    attributes["file_name"] = file_upload.name
+                case _:
+                    assert_never(file_upload)
+            with push_scope_optional(
+                self._logger, "upload_file", attributes=attributes
+            ):
+                if self._logger is not None:
+                    self._logger.log_info("Uploading file", attributes=attributes)
+                with file_upload_data(file_upload) as file_bytes:
+                    if file_bytes.bytes_data.read(1) == b"":
+                        raise UploadFailed(
+                            f"Failed to upload empty file: {file_bytes.name}"
+                        )
+                    file_bytes.bytes_data.seek(0)
+                    location = await aiotus.upload(
+                        creation_url,
+                        file_bytes.bytes_data,
+                        {"filename": file_bytes.name.encode()},
+                        client_session=session,
+                        config=aiotus.RetryConfiguration(
+                            ssl=not self._allow_insecure_tls
+                        ),
+                        chunksize=_CHUNK_SIZE,
+                    )
+                    if location is None:
+                        raise UploadFailed(f"Failed to upload: {file_bytes.name}")
+                    return UploadedFile(
+                        name=file_bytes.name, file_id=int(location.path.split("/")[-1])
+                    )
+
+    def upload_files(
+        self: Self, *, file_uploads: list[FileUpload]
+    ) -> list[UploadedFile]:
         return [
             asyncio.run(self._upload_file(file_upload)) for file_upload in file_uploads
         ]