From a4f067fb6b24f7e65e3dec726260d34ec5847fcd Mon Sep 17 00:00:00 2001 From: JordonPhillips Date: Fri, 20 Dec 2024 18:56:54 +0100 Subject: [PATCH] Always wrap event stream response bodies --- .../codegen/integration/RestJsonProtocolGenerator.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/integration/RestJsonProtocolGenerator.java b/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/integration/RestJsonProtocolGenerator.java index 5246a031b..e916980da 100644 --- a/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/integration/RestJsonProtocolGenerator.java +++ b/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/integration/RestJsonProtocolGenerator.java @@ -406,6 +406,7 @@ public void wrapEventStream(GenerationContext context, PythonWriter writer) { writer.addImports("aws_event_stream.aio", Set.of( "AWSDuplexEventStream", "AWSInputEventStream", "AWSOutputEventStream")); writer.addImport("smithy_json", "JSONCodec"); + writer.addImport("smithy_core.aio.types", "AsyncBytesReader"); writer.addImport("smithy_core.types", "TimestampFormat"); writer.addStdlibImport("typing", "Any"); @@ -417,7 +418,7 @@ public void wrapEventStream(GenerationContext context, PythonWriter writer) { payload_codec=codec, initial_response=operation_output, async_writer=execution_context.transport_request.body, # type: ignore - async_reader=execution_context.transport_response.body, # type: ignore + async_reader=AsyncBytesReader(execution_context.transport_response.body), # type: ignore deserializer=event_deserializer, # type: ignore ) else: @@ -430,7 +431,7 @@ public void wrapEventStream(GenerationContext context, PythonWriter writer) { return AWSOutputEventStream[Any, Any]( payload_codec=codec, initial_response=operation_output, - async_reader=execution_context.transport_response.body, # type: ignore + async_reader=AsyncBytesReader(execution_context.transport_response.body), # type: ignore deserializer=event_deserializer, # type: ignore ) """);