Cleanup

Author: nateprewitt

codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java

@@ -571,6 +571,9 @@ await sleep(retry_token.retry_delay)
         writer.pushState(new SignRequestSection());
         if (context.applicationProtocol().isHttpProtocol() && supportsAuth) {
             writer.addStdlibImport("re");
+            writer.addStdlibImport("typing", "Any");
+            writer.addImport("smithy_core.interfaces.identity", "Identity");
+            writer.addImport("smithy_core.types", "PropertyKey");
             writer.write("""
                             # Step 7i: sign the request
                             if auth_option and signer:
@@ -594,8 +597,11 @@ await sleep(retry_token.retry_delay)
                                 auth_value = fields["Authorization"].as_string()  # type: ignore
                                 signature = re.split("Signature=", auth_value)[-1]  # type: ignore
                                 context.properties["signature"] = signature.encode('utf-8')
-                                context.properties["identity"] = identity
-                                context.properties["signer_properties"] = auth_option.signer_properties
+
+                                identity_key = PropertyKey(key="identity", value_type=Identity | None)
+                                sp_key = PropertyKey(key="signer_properties", value_type=dict)
+                                context.properties[identity_key] = identity
+                                context.properties[sp_key] = auth_option.signer_properties
                     """);
         }
         writer.popState();

codegen/core/src/main/java/software/amazon/smithy/python/codegen/integrations/RestJsonProtocolGenerator.java

@@ -452,7 +452,7 @@ public void wrapOutputStream(GenerationContext context, PythonWriter writer) {
                                 transport_response.body  # type: ignore
                             ),
                             deserializer=event_deserializer,  # type: ignore
-                            signer=signer,  # type: ignore
                         )
                         """);
     }
+}

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

@@ -5,7 +5,7 @@
 from collections.abc import Iterator
 from contextlib import contextmanager
 from io import BytesIO
-from typing import Never, Protocol
+from typing import Never
 
 from smithy_core.codecs import Codec
 from smithy_core.schemas import Schema

packages/aws-event-stream/src/aws_event_stream/aio/__init__.py

@@ -14,9 +14,11 @@
 
 from .._private.serializers import EventSerializer as _EventSerializer
 from .._private.deserializers import EventDeserializer as _EventDeserializer
-from ..events import Event, EventMessage
+from ..events import Event, EventHeaderEncoder, EventMessage
 from ..exceptions import EventError
 
+from typing import Protocol
+
 logger = logging.getLogger(__name__)
 
 
@@ -57,8 +59,13 @@ async def send(self, event: E) -> None:
                 "Expected an event message to be serialized, but was None."
             )
         if self._signer is not None:
-            result = self._signer(result)
+            encoder = self._serializer.event_header_encoder_cls
+            result = await self._signer.sign_event(  # type: ignore
+                event_message=result,
+                event_encoder_cls=encoder,
+            )
 
+        assert isinstance(result, EventMessage)
         encoded_result = result.encode()
         try:
             logger.debug("Publishing serialized event: %s", result)

packages/aws-event-stream/src/aws_event_stream/events.py

@@ -387,7 +387,7 @@ def get_result(self) -> bytes:
             raise InvalidHeadersLength(len(result))
         return result
 
-    def encode_headers(self, headers: HEADERS_DICT):
+    def encode_headers(self, headers: HEADERS_DICT) -> None:
         """Encode a map of headers.
 
         :param headers: A mapping of headers to encode.

packages/aws-sdk-signers/src/aws_sdk_signers/signers.py

@@ -833,7 +833,7 @@ async def sign_event(
                 )
 
             timestamp = new_signing_properties["date"]
-            headers: dict[str, str | bytes] = {":date": date_obj}
+            headers: dict[str, str | bytes | datetime.datetime] = {":date": date_obj}
             encoder = event_encoder_cls()
             encoder.encode_headers(headers)
             encoded_headers = encoder.get_result()

packages/smithy-core/src/smithy_core/interfaces/identity.py

@@ -1,9 +1,10 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
 from datetime import datetime
-from typing import Protocol, TypedDict, TypeVar
+from typing import Protocol, TypedDict, TypeVar, runtime_checkable
 
 
+@runtime_checkable
 class Identity(Protocol):
     """An entity available to the client representing who the user is."""