Centralize event stream wrappers

Author: JordonPhillips

designs/event-streams.md

@@ -53,6 +53,8 @@ async with publisher:
     publisher.send(FooEvent(foo="bar"))
 ```
 
+Protocol implementations will be responsible for creating publishers.
+
 ## Event Receivers
 
 An `AsyncEventReceiver` is used to receive events from a service.
@@ -131,6 +133,8 @@ async for event in reciever:
     handle_event(event)
 ```
 
+Protocol implementations will be responsible for creating receivers.
+
 ### Errors
 
 Event streams may define modeled errors that may be sent over the stream. These
@@ -169,38 +173,30 @@ are handled by the following classes:
 * `OutputEventStream` is returned when the operation only has an output stream.
 
 ```python
-class DuplexEventStream[I: SerializableShape, O: DeserializableShape, R](Protocol):
-
-    input_stream: AsyncEventPublisher[I]
-
-    _output_stream: AsyncEventReceiver[O] | None = None
-    _response: R | None = None
-
-    @property
-    def output_stream(self) -> AsyncEventReceiver[O] | None:
-        return self._output_stream
-
-    @output_stream.setter
-    def output_stream(self, value: AsyncEventReceiver[O]) -> None:
-        self._output_stream = value
-
-    @property
-    def response(self) -> R | None:
-        return self._response
-
-    @response.setter
-    def response(self, value: R) -> None:
-        self._response = value
-
-    async def await_output(self) -> tuple[R, AsyncEventReceiver[O]]:
+class DuplexEventStream[
+    IE: SerializeableShape,
+    OE: DeserializeableShape,
+    O: DeserializeableShape,
+]:
+
+    input_stream: EventPublisher[IE]
+    output_stream: EventReceiver[OE] | None = None
+    response: O | None = None
+
+    def __init__(
+        self,
+        *,
+        input_stream: EventPublisher[IE],
+        output_future: Future[tuple[O, EventReceiver[OE]]],
+    ) -> None:
+        self.input_stream = input_stream
+        self._output_future = output_future
+
+    async def await_output(self) -> tuple[O, EventReceiver[OE]]:
         ...
 
     async def close(self) -> None:
-        if self.output_stream is None:
-            _, self.output_stream = await self.await_output()
-
-        await self.input_stream.close()
-        await self.output_stream.close()
+        ...
 
     async def __aenter__(self) -> Self:
         return self
@@ -209,21 +205,21 @@ class DuplexEventStream[I: SerializableShape, O: DeserializableShape, R](Protoco
         await self.close()
 
 
-class InputEventStream[I: SerializableShape, R](Protocol):
+class InputEventStream[IE: SerializeableShape, O]:
 
-    input_stream: AsyncEventPublisher[I]
+    input_stream: EventPublisher[IE]
+    response: O | None = None
 
-    _response: R | None = None
+    def __init__(
+        self,
+        *,
+        input_stream: EventPublisher[IE],
+        output_future: Future[O],
+    ) -> None:
+        self.input_stream = input_stream
+        self._output_future = output_future
 
-    @property
-    def response(self) -> R | None:
-        return self._response
-
-    @response.setter
-    def response(self, value: R) -> None:
-        self._response = value
-
-    async def await_output(self) -> R:
+    async def await_output(self) -> O:
         ...
 
     async def close(self) -> None:
@@ -236,11 +232,14 @@ class InputEventStream[I: SerializableShape, R](Protocol):
         await self.close()
 
 
-class OutputEventStream[O: DeserializableShape, R](Protocol):
+class OutputEventStream[OE: DeserializeableShape, O: DeserializeableShape]:
 
-    output_stream: AsyncEventReceiver[O]
-
-    response: R
+    output_stream: EventReceiver[OE]
+    response: O
+    
+    def __init__(self, output_stream: EventReceiver[OE], output: O) -> None:
+        self.output_stream = output_stream
+        self.response = output
 
     async def close(self) -> None:
         await self.output_stream.close()
@@ -290,6 +289,9 @@ with await client.output_operation() as stream:
         handle_event(event)
 ```
 
+All three output types are centrally located and will be constructed by filling
+in the relevant publishers and receivers from the protocol implementation.
+
 ## Event Structure
 
 Event messages are structurally similar to HTTP messages. They consist of a map

packages/aws-event-stream/pyproject.toml

@@ -6,7 +6,6 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "smithy-core",
-    "smithy-event-stream",
 ]
 
 [build-system]

packages/aws-event-stream/src/aws_event_stream/_private/deserializers.py

@@ -1,20 +1,16 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-import asyncio
 import datetime
 from collections.abc import Callable
 
-from smithy_core.aio.interfaces import AsyncByteStream
 from smithy_core.codecs import Codec
 from smithy_core.deserializers import (
-    DeserializeableShape,
     ShapeDeserializer,
     SpecificShapeDeserializer,
 )
 from smithy_core.schemas import Schema
 from smithy_core.shapes import ShapeType
 from smithy_core.utils import expect_type
-from smithy_event_stream.aio.interfaces import AsyncEventReceiver
 
 from ..events import HEADERS_DICT, Event
 from ..exceptions import EventError, UnmodeledEventError
@@ -28,59 +24,6 @@
 INITIAL_MESSAGE_TYPES = (INITIAL_REQUEST_EVENT_TYPE, INITIAL_RESPONSE_EVENT_TYPE)
 
 
-class AWSAsyncEventReceiver[E: DeserializeableShape](AsyncEventReceiver[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 IOError("Failed to read from stream.") from e
-            raise
-
-        if event is None:
-            return None
-
-        deserializer = EventDeserializer(
-            event=event,
-            payload_codec=self._payload_codec,
-            is_client_mode=self._is_client_mode,
-        )
-        result = self._deserializer(deserializer)
-        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
-
-
 class EventDeserializer(SpecificShapeDeserializer):
     def __init__(
         self, event: Event, payload_codec: Codec, is_client_mode: bool = True

packages/aws-event-stream/src/aws_event_stream/_private/serializers.py

@@ -1,24 +1,19 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-import asyncio
 import datetime
-from collections.abc import Callable, Iterator
+from collections.abc import Iterator
 from contextlib import contextmanager
 from io import BytesIO
 from typing import Never
 
-from smithy_core.aio.interfaces import AsyncWriter
 from smithy_core.codecs import Codec
-from smithy_core.exceptions import ExpectationNotMetException
 from smithy_core.schemas import Schema
 from smithy_core.serializers import (
     InterceptingSerializer,
-    SerializeableShape,
     ShapeSerializer,
     SpecificShapeSerializer,
 )
 from smithy_core.shapes import ShapeType
-from smithy_event_stream.aio.interfaces import AsyncEventPublisher
 
 from ..events import EventMessage, HEADER_VALUE, Short, Byte, Long
 from ..exceptions import InvalidHeaderValue
@@ -33,58 +28,6 @@
 _DEFAULT_BLOB_CONTENT_TYPE = "application/octet-stream"
 
 
-type Signer = Callable[[EventMessage], EventMessage]
-"""A function that takes an event message and signs it, and returns it signed."""
-
-
-class AWSAsyncEventPublisher[E: SerializeableShape](AsyncEventPublisher[E]):
-    def __init__(
-        self,
-        payload_codec: Codec,
-        async_writer: AsyncWriter,
-        signer: Signer | 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 IOError("Attempted to write to closed stream.")
-        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:
-            result = self._signer(result)
-
-        encoded_result = result.encode()
-        try:
-            await self._writer.write(encoded_result)
-        except Exception as e:
-            await self.close()
-            raise IOError("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 EventSerializer(SpecificShapeSerializer):
     def __init__(
         self,