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 ) """);