diff --git a/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/StructureGenerator.java b/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/StructureGenerator.java index 6703ab5f1..54275da9b 100644 --- a/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/StructureGenerator.java +++ b/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/StructureGenerator.java @@ -41,6 +41,7 @@ import software.amazon.smithy.model.traits.InputTrait; import software.amazon.smithy.model.traits.OutputTrait; import software.amazon.smithy.model.traits.SensitiveTrait; +import software.amazon.smithy.model.traits.StreamingTrait; import software.amazon.smithy.python.codegen.generators.MemberDeserializerGenerator; import software.amazon.smithy.python.codegen.generators.MemberSerializerGenerator; @@ -417,6 +418,9 @@ private void deserializeMembers(Collection members) { int index = 0; for (MemberShape member : members) { var target = model.expectShape(member.getTarget()); + if (target.hasTrait(StreamingTrait.class) && target.isUnionShape()) { + continue; + } writer.write(""" case $L: kwargs[$S] = ${C|} 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 e40e16ce3..39484cd84 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 @@ -314,10 +314,15 @@ protected void deserializePayloadBody( var symbolProvider = context.symbolProvider(); var memberName = symbolProvider.toMemberName(payloadBinding.getMember()); + var target = context.model().expectShape(payloadBinding.getMember().getTarget()); - if (payloadBinding.getMember().getMemberTrait(context.model(), StreamingTrait.class).isPresent()) { - writer.addImport("smithy_core.aio.types", "AsyncBytesReader"); - writer.write("kwargs[$S] = AsyncBytesReader(http_response.body)", memberName); + if (target.hasTrait(StreamingTrait.class)) { + // If the target is a union, nothing needs to be done here. The stream deserializer will + // take the response body directly and handle it. + if (target.isBlobShape()) { + writer.addImport("smithy_core.aio.types", "AsyncBytesReader"); + writer.write("kwargs[$S] = AsyncBytesReader(http_response.body)", memberName); + } return; } @@ -334,7 +339,6 @@ protected void deserializePayloadBody( """); } - var target = context.model().expectShape(payloadBinding.getMember().getTarget()); var deserializerSymbol = symbolProvider.toSymbol(target); writer.write("if body:").indent();