Serialize string/blob event payloads directly

Author: JordonPhillips

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

@@ -131,6 +131,8 @@ def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]:
 
         if (payload_member := self._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:
@@ -200,19 +202,30 @@ def write_timestamp(self, schema: "Schema", value: datetime.datetime) -> None:
         self._encoder.encode_timestamp(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_serializer: ShapeSerializer,
+        payload_struct_serializer: ShapeSerializer,
     ) -> None:
         self._header_serializer = header_serializer
-        self._payload_serializer = payload_serializer
+        self._payload_struct_serializer = payload_struct_serializer
 
     def before(self, schema: "Schema") -> ShapeSerializer:
         if EVENT_HEADER_TRAIT in schema.traits:
             return self._header_serializer
-        return self._payload_serializer
+        return self._payload_struct_serializer
 
     def after(self, schema: "Schema") -> None:
         pass

python-packages/aws-event-stream/tests/unit/_private/__init__.py

@@ -1,28 +1,27 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-from dataclasses import dataclass
 import datetime
-from typing import Any, Self, ClassVar, Literal
-
-from aws_event_stream.events import EventMessage, Short, Byte, Long
+from dataclasses import dataclass
+from typing import Any, ClassVar, Literal, Self
 
-from smithy_core.serializers import ShapeSerializer
 from smithy_core.deserializers import ShapeDeserializer
-from smithy_core.schemas import Schema
 from smithy_core.exceptions import SmithyException
-from smithy_core.shapes import ShapeID, ShapeType
-from smithy_core.traits import Trait
 from smithy_core.prelude import (
+    BLOB,
     BOOLEAN,
     BYTE,
-    SHORT,
     INTEGER,
     LONG,
-    BLOB,
+    SHORT,
     STRING,
     TIMESTAMP,
 )
+from smithy_core.schemas import Schema
+from smithy_core.serializers import ShapeSerializer
+from smithy_core.shapes import ShapeID, ShapeType
+from smithy_core.traits import Trait
 
+from aws_event_stream.events import Byte, EventMessage, Long, Short
 
 EVENT_HEADER_TRAIT = Trait(id=ShapeID("smithy.api#eventHeader"))
 EVENT_PAYLOAD_TRAIT = Trait(id=ShapeID("smithy.api#eventPayload"))
@@ -66,6 +65,22 @@
     },
 )
 
+SCHEMA_BLOB_PAYLOAD_EVENT = Schema.collection(
+    id=ShapeID("smithy.example#BlobPayloadEvent"),
+    members={
+        "header": {
+            "index": 0,
+            "target": STRING,
+            "traits": [EVENT_HEADER_TRAIT, REQUIRED_TRAIT],
+        },
+        "payload": {
+            "index": 1,
+            "target": BLOB,
+            "traits": [EVENT_PAYLOAD_TRAIT, REQUIRED_TRAIT],
+        },
+    },
+)
+
 SCHEMA_ERROR_EVENT = Schema.collection(
     id=ShapeID("smithy.example#ErrorEvent"),
     members={"message": {"index": 0, "target": STRING, "traits": [REQUIRED_TRAIT]}},
@@ -79,7 +94,8 @@
     members={
         "message": {"index": 0, "target": SCHEMA_MESSAGE_EVENT},
         "payload": {"index": 1, "target": SCHEMA_PAYLOAD_EVENT},
-        "error": {"index": 2, "target": SCHEMA_ERROR_EVENT},
+        "blobPayload": {"index": 2, "target": SCHEMA_BLOB_PAYLOAD_EVENT},
+        "error": {"index": 3, "target": SCHEMA_ERROR_EVENT},
     },
 )
 
@@ -273,6 +289,57 @@ def serialize_members(self, serializer: ShapeSerializer):
         serializer.write_struct(SCHEMA_EVENT_STREAM.members["payload"], self.value)
 
 
+@dataclass
+class BlobPayloadEvent:
+    header: str
+    payload: bytes
+
+    def serialize(self, serializer: ShapeSerializer):
+        with serializer.begin_struct(SCHEMA_BLOB_PAYLOAD_EVENT) as s:
+            self.serialize_members(s)
+
+    def serialize_members(self, serializer: ShapeSerializer) -> None:
+        serializer.write_string(
+            SCHEMA_BLOB_PAYLOAD_EVENT.members["header"], self.header
+        )
+        serializer.write_blob(
+            SCHEMA_BLOB_PAYLOAD_EVENT.members["payload"], self.payload
+        )
+
+    @classmethod
+    def deserialize(cls, deserializer: ShapeDeserializer) -> Self:
+        kwargs: dict[str, Any] = {}
+
+        def _consumer(schema: Schema, de: ShapeDeserializer) -> None:
+            match schema.expect_member_index():
+                case 0:
+                    kwargs["header"] = de.read_string(
+                        SCHEMA_BLOB_PAYLOAD_EVENT.members["header"]
+                    )
+                case 1:
+                    kwargs["payload"] = de.read_blob(
+                        SCHEMA_BLOB_PAYLOAD_EVENT.members["payload"]
+                    )
+                case _:
+                    raise SmithyException(f"Unexpected member schema: {schema}")
+
+        deserializer.read_struct(schema=SCHEMA_BLOB_PAYLOAD_EVENT, consumer=_consumer)
+        return cls(**kwargs)
+
+
+@dataclass
+class EventStreamBlobPayloadEvent:
+    value: BlobPayloadEvent
+
+    def serialize(self, serializer: ShapeSerializer):
+        serializer.write_struct(SCHEMA_EVENT_STREAM, self)
+
+    def serialize_members(self, serializer: ShapeSerializer):
+        serializer.write_struct(
+            SCHEMA_EVENT_STREAM.members["blobPayload"], self.value
+        )
+
+
 @dataclass
 class ErrorEvent:
     code: ClassVar[str] = "NoSuchResource"
@@ -326,7 +393,7 @@ def serialize_members(self, serializer: ShapeSerializer):
         raise SmithyException("Unknown union variants may not be serialized.")
 
 
-type EventStream = EventStreamMessageEvent | EventStreamPayloadEvent | EventStreamErrorEvent | EventStreamUnknownEvent
+type EventStream = EventStreamMessageEvent | EventStreamPayloadEvent | EventStreamBlobPayloadEvent | EventStreamErrorEvent | EventStreamUnknownEvent
 
 
 class EventStreamDeserializer:
@@ -350,6 +417,11 @@ def _consumer(self, schema: Schema, de: ShapeDeserializer) -> None:
                 self._set_result(EventStreamPayloadEvent(PayloadEvent.deserialize(de)))
 
             case 2:
+                self._set_result(
+                    EventStreamBlobPayloadEvent(BlobPayloadEvent.deserialize(de))
+                )
+
+            case 3:
                 self._set_result(EventStreamErrorEvent(ErrorEvent.deserialize(de)))
 
             case _:
@@ -528,7 +600,21 @@ def _consumer(schema: Schema, de: ShapeDeserializer) -> None:
                 "header": "header",
                 ":content-type": "text/plain",
             },
-            payload=b'"payload"',
+            payload=b"payload",
+        ),
+    ),
+    (
+        EventStreamBlobPayloadEvent(
+            BlobPayloadEvent(header="header", payload=b"\x07beep\x07")
+        ),
+        EventMessage(
+            headers={
+                ":message-type": "event",
+                ":event-type": "blobPayload",
+                "header": "header",
+                ":content-type": "application/octet-stream",
+            },
+            payload=b"\x07beep\x07",
         ),
     ),
     (