smithy-aws-event-stream 0.0.1__tar.gz

smithy_aws_event_stream-0.0.1/.gitignore

@@ -0,0 +1,36 @@
+# Eclipse
+.classpath
+.project
+.settings/
+
+# Intellij
+.idea/
+*.iml
+*.iws
+
+# VS Code
+.vscode/
+
+# Mac
+.DS_Store
+
+# Maven
+target/
+**/dependency-reduced-pom.xml
+
+# Gradle
+**/.gradle
+build/
+*/out/
+**/*/out/
+
+# Python
+__pycache__
+*.swp
+.mypy_cache
+.coverage
+coverage.xml
+htmlcov
+*.egg-info
+dist
+.venv

smithy_aws_event_stream-0.0.1/CHANGES.md

@@ -0,0 +1,11 @@
+# Changelog
+
+## Unreleased
+
+* <Add new items here>
+
+## v0.0.1
+
+### Features
+* Added basic support for the [Amazon Event Stream](https://smithy.io/2.0/aws/amazon-eventstream.html) specification.
+* Added support for unidirectdional and bidirectional event streams.

smithy_aws_event_stream-0.0.1/NOTICE

@@ -0,0 +1 @@
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

smithy_aws_event_stream-0.0.1/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.4
+Name: smithy-aws-event-stream
+Version: 0.0.1
+Summary: Smithy components for Amazon Event Streams.
+License-File: NOTICE
+Requires-Python: >=3.12
+Requires-Dist: smithy-core

smithy_aws_event_stream-0.0.1/pyproject.toml

@@ -0,0 +1,21 @@
+[project]
+name = "smithy-aws-event-stream"
+version = "0.0.1"
+description = "Smithy components for Amazon Event Streams."
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "smithy-core",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build]
+exclude = [
+  "tests",
+]
+
+[tool.ruff]
+src = ["src"]

smithy_aws_event_stream-0.0.1/src/smithy_aws_event_stream/__init__.py

@@ -0,0 +1,6 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib.metadata
+
+__version__: str = importlib.metadata.version("smithy-aws-event-stream")

smithy_aws_event_stream-0.0.1/src/smithy_aws_event_stream/_private/__init__.py

@@ -0,0 +1,14 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from smithy_core.schemas import Schema
+from smithy_core.traits import EventPayloadTrait
+
+INITIAL_REQUEST_EVENT_TYPE = "initial-request"
+INITIAL_RESPONSE_EVENT_TYPE = "initial-response"
+
+
+def get_payload_member(schema: Schema) -> Schema | None:
+    for member in schema.members.values():
+        if EventPayloadTrait in member:
+            return member
+    return None

smithy_aws_event_stream-0.0.1/src/smithy_aws_event_stream/_private/deserializers.py

@@ -0,0 +1,166 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import datetime
+import logging
+from collections.abc import Callable
+
+from smithy_core.codecs import Codec
+from smithy_core.deserializers import (
+    ShapeDeserializer,
+    SpecificShapeDeserializer,
+)
+from smithy_core.schemas import Schema
+from smithy_core.shapes import ShapeType
+from smithy_core.traits import EventHeaderTrait
+from smithy_core.utils import expect_type
+
+from ..events import HEADERS_DICT, Event
+from ..exceptions import EventError, UnmodeledEventError
+from . import (
+    INITIAL_REQUEST_EVENT_TYPE,
+    INITIAL_RESPONSE_EVENT_TYPE,
+    get_payload_member,
+)
+
+logger = logging.getLogger(__name__)
+
+INITIAL_MESSAGE_TYPES = (INITIAL_REQUEST_EVENT_TYPE, INITIAL_RESPONSE_EVENT_TYPE)
+
+
+class EventDeserializer(SpecificShapeDeserializer):
+    def __init__(
+        self, event: Event, payload_codec: Codec, is_client_mode: bool = True
+    ) -> None:
+        self._event = event
+        self._payload_codec = payload_codec
+        self._is_client_mode = is_client_mode
+
+    def read_struct(
+        self,
+        schema: Schema,
+        consumer: Callable[[Schema, ShapeDeserializer], None],
+    ) -> None:
+        headers = self._event.message.headers
+
+        match headers.get(":message-type"):
+            case "event":
+                member_name = expect_type(str, headers[":event-type"])
+                if member_name in INITIAL_MESSAGE_TYPES:
+                    # If it's an initial message, skip straight to deserialization.
+                    message_deserializer = self._create_deserializer(schema, headers)
+                    message_deserializer.read_struct(schema, consumer)
+                else:
+                    member_schema = schema.members[member_name]
+                    message_deserializer = self._create_deserializer(
+                        member_schema, headers
+                    )
+                    consumer(member_schema, message_deserializer)
+            case "exception":
+                member_name = expect_type(str, headers[":exception-type"])
+                member_schema = schema.members[member_name]
+                message_deserializer = self._create_deserializer(member_schema, headers)
+                consumer(member_schema, message_deserializer)
+            case "error":
+                # The `application/vnd.amazon.eventstream` format allows for explicitly
+                # unmodeled exceptions. These exceptions MUST have the `:error-code`
+                # and `:error-message` headers set, and they MUST be strings.
+                raise UnmodeledEventError(
+                    expect_type(str, headers[":error-code"]),
+                    expect_type(str, headers[":error-message"]),
+                )
+            case _:
+                raise EventError(f"Unknown event structure: {self._event}")
+
+    def _create_deserializer(
+        self, schema: Schema, headers: HEADERS_DICT
+    ) -> ShapeDeserializer:
+        payload_member = get_payload_member(schema)
+        payload_deserializer = self._create_payload_deserializer(payload_member)
+        return EventMessageDeserializer(headers, payload_deserializer, payload_member)
+
+    def _create_payload_deserializer(
+        self, payload_member: Schema | None
+    ) -> ShapeDeserializer | None:
+        if not self._event.message.payload:
+            return
+
+        if payload_member is not None and payload_member.shape_type in (
+            ShapeType.BLOB,
+            ShapeType.STRING,
+        ):
+            return RawPayloadDeserializer(self._event.message.payload)
+
+        return self._payload_codec.create_deserializer(self._event.message.payload)
+
+
+class RawPayloadDeserializer(SpecificShapeDeserializer):
+    def __init__(self, payload: bytes) -> None:
+        self._payload = payload
+
+    def read_string(self, schema: Schema) -> str:
+        return self._payload.decode("utf-8")
+
+    def read_blob(self, schema: Schema) -> bytes:
+        return self._payload
+
+
+class EventMessageDeserializer(SpecificShapeDeserializer):
+    def __init__(
+        self,
+        headers: HEADERS_DICT,
+        payload_deserializer: ShapeDeserializer | None,
+        payload_member: Schema | None,
+    ) -> None:
+        self._headers = headers
+        self._payload_deserializer = payload_deserializer
+        self._payload_member = payload_member
+
+    def read_struct(
+        self,
+        schema: Schema,
+        consumer: Callable[[Schema, ShapeDeserializer], None],
+    ) -> None:
+        headers_deserializer = EventHeaderDeserializer(self._headers)
+        for key in self._headers.keys():
+            member_schema = schema.members.get(key)
+            if member_schema is not None and EventHeaderTrait in member_schema:
+                consumer(member_schema, headers_deserializer)
+
+        if self._payload_deserializer:
+            if self._payload_member is not None:
+                consumer(self._payload_member, self._payload_deserializer)
+            else:
+                self._payload_deserializer.read_struct(schema, consumer)
+
+
+class EventHeaderDeserializer(SpecificShapeDeserializer):
+    def __init__(self, headers: HEADERS_DICT) -> None:
+        self._headers = headers
+
+    def read_boolean(self, schema: "Schema") -> bool:
+        return expect_type(bool, self._headers[schema.expect_member_name()])
+
+    def read_blob(self, schema: "Schema") -> bytes:
+        return expect_type(bytes, self._headers[schema.expect_member_name()])
+
+    def read_byte(self, schema: "Schema") -> int:
+        return expect_type(int, self._headers[schema.expect_member_name()])
+
+    def read_short(self, schema: "Schema") -> int:
+        return expect_type(int, self._headers[schema.expect_member_name()])
+
+    def read_integer(self, schema: "Schema") -> int:
+        return expect_type(int, self._headers[schema.expect_member_name()])
+
+    def read_long(self, schema: "Schema") -> int:
+        return expect_type(int, self._headers[schema.expect_member_name()])
+
+    def read_string(self, schema: "Schema") -> str:
+        return expect_type(str, self._headers[schema.expect_member_name()])
+
+    def read_timestamp(self, schema: "Schema") -> datetime.datetime:
+        # TODO: do we support timestamp format here? One would assume not since the
+        # format has a specific timestamp type.
+        return expect_type(
+            datetime.datetime, self._headers[schema.expect_member_name()]
+        )

smithy_aws_event_stream-0.0.1/src/smithy_aws_event_stream/_private/serializers.py

@@ -0,0 +1,181 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import datetime
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
+from io import BytesIO
+from typing import Never
+
+from smithy_core.codecs import Codec
+from smithy_core.schemas import Schema
+from smithy_core.serializers import (
+    InterceptingSerializer,
+    ShapeSerializer,
+    SpecificShapeSerializer,
+)
+from smithy_core.shapes import ShapeType
+from smithy_core.traits import ErrorTrait, EventHeaderTrait, MediaTypeTrait
+
+from ..events import HEADER_VALUE, Byte, EventHeaderEncoder, EventMessage, Long, Short
+from ..exceptions import InvalidHeaderValue
+from . import (
+    INITIAL_REQUEST_EVENT_TYPE,
+    INITIAL_RESPONSE_EVENT_TYPE,
+    get_payload_member,
+)
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_STRING_CONTENT_TYPE = "text/plain"
+_DEFAULT_BLOB_CONTENT_TYPE = "application/octet-stream"
+
+
+class EventSerializer(SpecificShapeSerializer):
+    def __init__(
+        self,
+        payload_codec: Codec,
+        is_client_mode: bool = True,
+    ) -> None:
+        self._payload_codec = payload_codec
+        self._result: EventMessage | None = None
+        if is_client_mode:
+            self._initial_message_event_type = INITIAL_REQUEST_EVENT_TYPE
+        else:
+            self._initial_message_event_type = INITIAL_RESPONSE_EVENT_TYPE
+        self.event_header_encoder_cls = EventHeaderEncoder
+
+    def get_result(self) -> EventMessage | None:
+        return self._result
+
+    @contextmanager
+    def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]:
+        # Event stream definitions are unions. Nothing about the union shape actually
+        # matters for the purposes of event stream serialization though, we just care
+        # about the specific member we're serializing. So here we yield immediately,
+        # and the next time this method is called it'll be by the member that we
+        # actually care about.
+        #
+        # Note that if we're serializing an operation input or output, it won't be a
+        # union at all, so this won't get triggered. Thankfully, that's what we want.
+        if schema.shape_type is ShapeType.UNION:
+            try:
+                yield self
+            finally:
+                return
+
+        headers: dict[str, HEADER_VALUE] = {}
+
+        if ErrorTrait in schema:
+            headers[":message-type"] = "exception"
+            headers[":exception-type"] = schema.expect_member_name()
+        else:
+            headers[":message-type"] = "event"
+            if schema.member_name is None:
+                # If there's no member name, that must mean that the structure is
+                # either an input or output structure, and so this represents the
+                # initial message.
+                headers[":event-type"] = self._initial_message_event_type
+            else:
+                headers[":event-type"] = schema.member_name
+
+        payload = BytesIO()
+        payload_serializer: ShapeSerializer = self._payload_codec.create_serializer(
+            payload
+        )
+
+        header_serializer = EventHeaderSerializer(headers)
+
+        media_type = self._payload_codec.media_type
+
+        if (payload_member := get_payload_member(schema)) is not None:
+            media_type = self._get_payload_media_type(payload_member, media_type)
+            if payload_member.shape_type in (ShapeType.BLOB, ShapeType.STRING):
+                payload_serializer = RawPayloadSerializer(payload)
+            yield EventStreamBindingSerializer(header_serializer, payload_serializer)
+        else:
+            with payload_serializer.begin_struct(schema) as body_serializer:
+                yield EventStreamBindingSerializer(header_serializer, body_serializer)
+
+        payload_bytes = payload.getvalue()
+        if payload_bytes:
+            headers[":content-type"] = media_type
+
+        self._result = EventMessage(headers=headers, payload=payload_bytes)
+
+    def _get_payload_media_type(self, schema: Schema, default: str) -> str:
+        if (media_type := schema.get_trait(MediaTypeTrait)) is not None:
+            return media_type.value
+
+        match schema.shape_type:
+            case ShapeType.STRING:
+                return _DEFAULT_STRING_CONTENT_TYPE
+            case ShapeType.BLOB:
+                return _DEFAULT_BLOB_CONTENT_TYPE
+            case _:
+                return default
+
+
+class EventHeaderSerializer(SpecificShapeSerializer):
+    def __init__(self, headers: dict[str, HEADER_VALUE]) -> None:
+        self._headers = headers
+
+    def _invalid_state(
+        self, schema: "Schema | None" = None, message: str | None = None
+    ) -> Never:
+        if message is None:
+            message = f"Invalid header value type: {schema}"
+        raise InvalidHeaderValue(message)
+
+    def write_boolean(self, schema: "Schema", value: bool) -> None:
+        self._headers[schema.expect_member_name()] = value
+
+    def write_byte(self, schema: "Schema", value: int) -> None:
+        self._headers[schema.expect_member_name()] = Byte(value)
+
+    def write_short(self, schema: "Schema", value: int) -> None:
+        self._headers[schema.expect_member_name()] = Short(value)
+
+    def write_integer(self, schema: "Schema", value: int) -> None:
+        self._headers[schema.expect_member_name()] = value
+
+    def write_long(self, schema: "Schema", value: int) -> None:
+        self._headers[schema.expect_member_name()] = Long(value)
+
+    def write_string(self, schema: "Schema", value: str) -> None:
+        self._headers[schema.expect_member_name()] = value
+
+    def write_blob(self, schema: "Schema", value: bytes) -> None:
+        self._headers[schema.expect_member_name()] = value
+
+    def write_timestamp(self, schema: "Schema", value: datetime.datetime) -> None:
+        self._headers[schema.expect_member_name()] = value
+
+
+class RawPayloadSerializer(SpecificShapeSerializer):
+    def __init__(self, payload: BytesIO) -> None:
+        self._payload = payload
+
+    def write_string(self, schema: "Schema", value: str) -> None:
+        self._payload.write(value.encode("utf-8"))
+
+    def write_blob(self, schema: "Schema", value: bytes) -> None:
+        self._payload.write(value)
+
+
+class EventStreamBindingSerializer(InterceptingSerializer):
+    def __init__(
+        self,
+        header_serializer: EventHeaderSerializer,
+        payload_struct_serializer: ShapeSerializer,
+    ) -> None:
+        self._header_serializer = header_serializer
+        self._payload_struct_serializer = payload_struct_serializer
+
+    def before(self, schema: "Schema") -> ShapeSerializer:
+        if EventHeaderTrait in schema:
+            return self._header_serializer
+        return self._payload_struct_serializer
+
+    def after(self, schema: "Schema") -> None:
+        pass

smithy_aws_event_stream-0.0.1/src/smithy_aws_event_stream/aio/__init__.py

@@ -0,0 +1,141 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import asyncio
+import logging
+from collections.abc import Callable
+from typing import Protocol
+
+from smithy_core.aio.interfaces import AsyncByteStream, AsyncWriter
+from smithy_core.aio.interfaces.eventstream import EventPublisher, EventReceiver
+from smithy_core.codecs import Codec
+from smithy_core.deserializers import DeserializeableShape, ShapeDeserializer
+from smithy_core.exceptions import ExpectationNotMetException
+from smithy_core.serializers import SerializeableShape
+
+from .._private.deserializers import EventDeserializer as _EventDeserializer
+from .._private.serializers import EventSerializer as _EventSerializer
+from ..events import Event, EventHeaderEncoder, EventMessage
+from ..exceptions import EventError
+
+logger = logging.getLogger(__name__)
+
+
+class EventSigner(Protocol):
+    """A signer to manage credentials and EventMessages for an Event Stream lifecyle."""
+
+    async def sign_event(
+        self,
+        *,
+        event_message: EventMessage,
+        event_encoder_cls: type[EventHeaderEncoder],
+    ) -> EventMessage: ...
+
+
+class AWSEventPublisher[E: SerializeableShape](EventPublisher[E]):
+    def __init__(
+        self,
+        payload_codec: Codec,
+        async_writer: AsyncWriter,
+        signer: EventSigner | None = None,
+        is_client_mode: bool = True,
+    ):
+        self._writer = async_writer
+        self._signer = signer
+        self._serializer = _EventSerializer(
+            payload_codec=payload_codec, is_client_mode=is_client_mode
+        )
+        self._closed = False
+
+    async def send(self, event: E) -> None:
+        if self._closed:
+            raise OSError("Attempted to write to closed stream.")
+        logger.debug("Preparing to publish event: %s", event)
+        event.serialize(self._serializer)
+        result = self._serializer.get_result()
+        if result is None:
+            raise ExpectationNotMetException(
+                "Expected an event message to be serialized, but was None."
+            )
+        if self._signer is not None:
+            encoder = self._serializer.event_header_encoder_cls
+            result = await self._signer.sign_event(
+                event_message=result,
+                event_encoder_cls=encoder,
+            )
+
+        encoded_result = result.encode()
+        try:
+            logger.debug("Publishing serialized event: %s", result)
+            await self._writer.write(encoded_result)
+        except Exception as e:
+            await self.close()
+            raise OSError("Failed to write to stream.") from e
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+
+        if (close := getattr(self._writer, "close", None)) is not None:
+            if asyncio.iscoroutine(result := close()):
+                await result
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+
+class AWSEventReceiver[E: DeserializeableShape](EventReceiver[E]):
+    def __init__(
+        self,
+        payload_codec: Codec,
+        source: AsyncByteStream,
+        deserializer: Callable[[ShapeDeserializer], E],
+        is_client_mode: bool = True,
+    ) -> None:
+        self._payload_codec = payload_codec
+        self._source = source
+        self._is_client_mode = is_client_mode
+        self._deserializer = deserializer
+        self._closed = False
+
+    async def receive(self) -> E | None:
+        if self._closed:
+            return None
+
+        try:
+            event = await Event.decode_async(self._source)
+        except Exception as e:
+            await self.close()
+            if not isinstance(e, EventError):
+                raise OSError("Failed to read from stream.") from e
+            raise
+
+        if event is None:
+            logger.debug("No event received from the source.")
+            return None
+        logger.debug("Received raw event: %s", event)
+
+        deserializer = _EventDeserializer(
+            event=event,
+            payload_codec=self._payload_codec,
+            is_client_mode=self._is_client_mode,
+        )
+        result = self._deserializer(deserializer)
+        logger.debug("Successfully deserialized event: %s", result)
+        if isinstance(getattr(result, "value"), Exception):
+            raise result.value  # type: ignore
+        return result
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+
+        if (close := getattr(self._source, "close", None)) is not None:
+            if asyncio.iscoroutine(result := close()):
+                await result
+
+    @property
+    def closed(self) -> bool:
+        return self._closed