Update generated stream wrappers

Author: JordonPhillips

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

@@ -200,6 +200,13 @@ def _classify_error(
         if (hasStreaming) {
             writer.addStdlibImports("typing", Set.of("Any", "Awaitable"));
             writer.addStdlibImport("asyncio");
+
+            writer.addImports("smithy_core.aio.eventstream",
+                    Set.of(
+                            "InputEventStream",
+                            "OutputEventStream",
+                            "DuplexEventStream"));
+            writer.addImport("smithy_core.aio.interfaces.eventstream", "EventReceiver");
             writer.write(
                     """
                             async def _input_stream[Input: SerializeableShape, Output: DeserializeableShape](
@@ -218,6 +225,10 @@ def _classify_error(
                                 ))
                                 request_context = await request_future
                                 ${5C|}
+                                return InputEventStream[Any, Any](
+                                    input_stream=publisher,
+                                    output_future=awaitable_output,
+                                )
 
                             async def _output_stream[Input: SerializeableShape, Output: DeserializeableShape](
                                 self,
@@ -236,6 +247,10 @@ def _classify_error(
                                 )
                                 transport_response = await response_future
                                 ${6C|}
+                                return OutputEventStream[Any, Any](
+                                    output_stream=receiver,
+                                    output=output
+                                )
 
                             async def _duplex_stream[Input: SerializeableShape, Output: DeserializeableShape](
                                 self,
@@ -255,15 +270,34 @@ def _classify_error(
                                     response_future=response_future
                                 ))
                                 request_context = await request_future
-                                ${7C|}
+                                ${5C|}
+                                output_future = asyncio.create_task(self._wrap_duplex_output(
+                                    response_future, awaitable_output, config, operation_name,
+                                    event_deserializer
+                                ))
+                                return DuplexEventStream[Any, Any, Any](
+                                    input_stream=publisher,
+                                    output_future=output_future,
+                                )
+
+                            async def _wrap_duplex_output(
+                                self,
+                                response_future: Future[$3T],
+                                awaitable_output: Future[Any],
+                                config: $4T,
+                                operation_name: str,
+                                event_deserializer: Callable[[ShapeDeserializer], Any],
+                            ) -> tuple[Any, EventReceiver[Any]]:
+                                transport_response = await response_future
+                                ${6C|}
+                                return await awaitable_output, receiver
                             """,
                     pluginSymbol,
                     transportRequest,
                     transportResponse,
                     configSymbol,
                     writer.consumer(w -> context.protocolGenerator().wrapInputStream(context, w)),
-                    writer.consumer(w -> context.protocolGenerator().wrapOutputStream(context, w)),
-                    writer.consumer(w -> context.protocolGenerator().wrapDuplexStream(context, w)));
+                    writer.consumer(w -> context.protocolGenerator().wrapOutputStream(context, w)));
         }
         writer.addStdlibImport("typing", "Any");
         writer.addStdlibImport("asyncio", "iscoroutine");
@@ -872,7 +906,6 @@ private void generateEventStreamOperation(PythonWriter writer, OperationShape op
 
         if (inputStreamSymbol != null) {
             if (outputStreamSymbol != null) {
-                writer.addImport("smithy_event_stream.aio.interfaces", "DuplexEventStream");
                 writer.write("""
                         async def ${operationName:L}(
                             self,
@@ -922,7 +955,6 @@ raise NotImplementedError()
                         """, writer.consumer(w -> writeSharedOperationInit(w, operation, input)));
             }
         } else {
-            writer.addImport("smithy_event_stream.aio.interfaces", "OutputEventStream");
             writer.write("""
                     async def ${operationName:L}(
                         self,

codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ProtocolGenerator.java

@@ -157,6 +157,4 @@ default void generateProtocolTests(GenerationContext context) {}
     default void wrapInputStream(GenerationContext context, PythonWriter writer) {}
 
     default void wrapOutputStream(GenerationContext context, PythonWriter writer) {}
-
-    default void wrapDuplexStream(GenerationContext context, PythonWriter writer) {}
 }

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

@@ -396,13 +396,12 @@ public void wrapInputStream(GenerationContext context, PythonWriter writer) {
         writer.addImport("smithy_json", "JSONCodec");
         writer.addImport("smithy_core.aio.types", "AsyncBytesReader");
         writer.addImport("smithy_core.types", "TimestampFormat");
-        writer.addImport("aws_event_stream.aio", "AWSInputEventStream");
+        writer.addImport("aws_event_stream.aio", "AWSEventPublisher");
         writer.write(
                 """
                         codec = JSONCodec(default_timestamp_format=TimestampFormat.EPOCH_SECONDS)
-                        return AWSInputEventStream[Any, Any](
+                        publisher = AWSEventPublisher[Any](
                             payload_codec=codec,
-                            awaitable_output=awaitable_output,
                             async_writer=request_context.transport_request.body,  # type: ignore
                         )
                         """);
@@ -415,39 +414,17 @@ public void wrapOutputStream(GenerationContext context, PythonWriter writer) {
         writer.addImport("smithy_json", "JSONCodec");
         writer.addImport("smithy_core.aio.types", "AsyncBytesReader");
         writer.addImport("smithy_core.types", "TimestampFormat");
-        writer.addImport("aws_event_stream.aio", "AWSOutputEventStream");
+        writer.addImport("aws_event_stream.aio", "AWSEventReceiver");
         writer.write(
                 """
                         codec = JSONCodec(default_timestamp_format=TimestampFormat.EPOCH_SECONDS)
-                        return AWSOutputEventStream[Any, Any](
+                        receiver = AWSEventReceiver(
                             payload_codec=codec,
-                            initial_response=output,
-                            async_reader=AsyncBytesReader(
+                            source=AsyncBytesReader(
                                 transport_response.body  # type: ignore
                             ),
                             deserializer=event_deserializer,  # type: ignore
                         )
                         """);
     }
-
-    @Override
-    public void wrapDuplexStream(GenerationContext context, PythonWriter writer) {
-        writer.addDependency(SmithyPythonDependency.SMITHY_JSON);
-        writer.addDependency(SmithyPythonDependency.AWS_EVENT_STREAM);
-        writer.addImport("smithy_json", "JSONCodec");
-        writer.addImport("smithy_core.aio.types", "AsyncBytesReader");
-        writer.addImport("smithy_core.types", "TimestampFormat");
-        writer.addImport("aws_event_stream.aio", "AWSDuplexEventStream");
-        writer.write(
-                """
-                        codec = JSONCodec(default_timestamp_format=TimestampFormat.EPOCH_SECONDS)
-                        return AWSDuplexEventStream[Any, Any, Any](
-                            payload_codec=codec,
-                            async_writer=request_context.transport_request.body,  # type: ignore
-                            awaitable_output=awaitable_output,
-                            awaitable_response=response_future,
-                            deserializer=event_deserializer,  # type: ignore
-                        )
-                        """);
-    }
 }