Implement event stream methods fro restJson1

Author: JordonPhillips

packages/smithy-aws-core/src/smithy_aws_core/aio/protocols.py

@@ -1,19 +1,59 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-from typing import Any, Final
-
+from collections.abc import Callable
+from inspect import iscoroutinefunction
+from typing import TYPE_CHECKING, Any, Final
+
+from smithy_core.aio.interfaces import AsyncWriter
+from smithy_core.aio.interfaces.auth import AuthScheme
+from smithy_core.aio.interfaces.eventstream import EventPublisher, EventReceiver
+from smithy_core.aio.types import AsyncBytesReader
 from smithy_core.codecs import Codec
-from smithy_core.exceptions import DiscriminatorError
+from smithy_core.deserializers import DeserializeableShape, ShapeDeserializer
+from smithy_core.exceptions import (
+    DiscriminatorError,
+    MissingDependencyError,
+    UnsupportedStreamError,
+)
+from smithy_core.interfaces import TypedProperties
 from smithy_core.schemas import APIOperation, Schema
+from smithy_core.serializers import SerializeableShape
 from smithy_core.shapes import ShapeID, ShapeType
 from smithy_core.types import TimestampFormat
-from smithy_http.aio.interfaces import HTTPErrorIdentifier, HTTPResponse
+from smithy_http.aio.interfaces import HTTPErrorIdentifier, HTTPRequest, HTTPResponse
 from smithy_http.aio.protocols import HttpBindingClientProtocol
 from smithy_json import JSONCodec, JSONDocument
 
 from ..traits import RestJson1Trait
 from ..utils import parse_document_discriminator, parse_error_code
 
+try:
+    from smithy_aws_event_stream.aio import (
+        AWSEventPublisher,
+        AWSEventReceiver,
+        SigningConfig,
+    )
+
+    _HAS_EVENT_STREAM = True
+except ImportError:
+    _HAS_EVENT_STREAM = False  # type: ignore
+
+if TYPE_CHECKING:
+    from smithy_aws_event_stream.aio import (
+        AWSEventPublisher,
+        AWSEventReceiver,
+        SigningConfig,
+    )
+    from typing_extensions import TypeForm
+
+
+def _assert_event_stream_capable() -> None:
+    if not _HAS_EVENT_STREAM:
+        raise MissingDependencyError(
+            "Attempted to use event streams, but smithy-aws-event-stream "
+            "is not installed."
+        )
+
 
 class AWSErrorIdentifier(HTTPErrorIdentifier):
     _HEADER_KEY: Final = "x-amzn-errortype"
@@ -80,3 +120,67 @@ def content_type(self) -> str:
     @property
     def error_identifier(self) -> HTTPErrorIdentifier:
         return self._error_identifier
+
+    def create_event_publisher[
+        OperationInput: SerializeableShape,
+        OperationOutput: DeserializeableShape,
+        Event: SerializeableShape,
+    ](
+        self,
+        *,
+        operation: "APIOperation[OperationInput, OperationOutput]",
+        request: HTTPRequest,
+        event_type: "TypeForm[Event]",
+        context: TypedProperties,
+        auth_scheme: AuthScheme[Any, Any, Any, Any] | None = None,
+    ) -> EventPublisher[Event]:
+        _assert_event_stream_capable()
+        signing_config: SigningConfig | None = None
+        if auth_scheme is not None:
+            event_signer = auth_scheme.event_signer(request=request)
+            if event_signer is not None:
+                signing_config = SigningConfig(
+                    signer=event_signer,
+                    signing_properties=auth_scheme.signer_properties(context=context),
+                    identity_resolver=auth_scheme.identity_resolver(context=context),
+                    identity_properties=auth_scheme.identity_properties(
+                        context=context
+                    ),
+                )
+
+        # The HTTP body must be an async writeable. The HTTP client bindings are
+        # responsible for ensuring this is the case. The CRT bindings, for example,
+        # will set the body to an instance of BufferableByteStream.
+        body = request.body
+        if not isinstance(body, AsyncWriter) or not iscoroutinefunction(body.write):
+            raise UnsupportedStreamError(
+                "Input streams require an async write function, but none was present "
+                "on the serialized HTTP request."
+            )
+
+        return AWSEventPublisher[Event](
+            payload_codec=self.payload_codec,
+            async_writer=body,
+            signing_config=signing_config,
+        )
+
+    def create_event_receiver[
+        OperationInput: SerializeableShape,
+        OperationOutput: DeserializeableShape,
+        Event: DeserializeableShape,
+    ](
+        self,
+        *,
+        operation: "APIOperation[OperationInput, OperationOutput]",
+        request: HTTPRequest,
+        response: HTTPResponse,
+        event_type: "TypeForm[Event]",
+        event_deserializer: Callable[[ShapeDeserializer], Event],
+        context: TypedProperties,
+    ) -> EventReceiver[Event]:
+        _assert_event_stream_capable()
+        return AWSEventReceiver(
+            payload_codec=self.payload_codec,
+            source=AsyncBytesReader(response.body),
+            deserializer=event_deserializer,
+        )