From 965955b45c8891f528a69a89dcacfd9948866f17 Mon Sep 17 00:00:00 2001 From: Joe Wu Date: Thu, 31 Jul 2025 14:10:03 -0700 Subject: [PATCH] Fix trait codegen for structure's string member with @idRef --- .../smithy/traitcodegen/test/CreatesTraitTest.java | 7 +++++++ .../traitcodegen/test/LoadsFromModelTest.java | 8 ++++++++ .../struct-with-idref-member-trait.smithy | 12 ++++++++++++ .../traitcodegen/generators/BuilderGenerator.java | 5 +++++ .../traitcodegen/generators/FromNodeGenerator.java | 9 +++++++++ .../smithy/traitcodegen/TraitCodegenPluginTest.java | 2 +- .../src/test/resources/META-INF/smithy/manifest | 1 + .../struct-with-idref-member-trait.smithy | 13 +++++++++++++ 8 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 smithy-trait-codegen/src/it/resources/software/amazon/smithy/traitcodegen/test/structures/struct-with-idref-member-trait.smithy create mode 100644 smithy-trait-codegen/src/test/resources/META-INF/smithy/structures/struct-with-idref-member-trait.smithy diff --git a/smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java b/smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java index d50282c3996..63584c77f8d 100644 --- a/smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java +++ b/smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java @@ -41,6 +41,7 @@ import com.example.traits.structures.BasicAnnotationTrait; import com.example.traits.structures.NestedA; import com.example.traits.structures.NestedB; +import com.example.traits.structures.StructWithIdrefMemberTrait; import com.example.traits.structures.StructWithListOfMapTrait; import com.example.traits.structures.StructWithUniqueItemsListTrait; import com.example.traits.structures.StructureTrait; @@ -200,6 +201,12 @@ static Stream createTraitTests() { .addItems(SetUtils.of("a", "b", "c")) .build() .toNode()), + Arguments.of(StructWithIdrefMemberTrait.ID, + StructWithIdrefMemberTrait.builder() + .idRefMemberA(ShapeId.from("test.smithy.traitcodegen#a")) + .idRefMemberB(ShapeId.from("test.smithy.traitcodegen#b")) + .build() + .toNode()), // Timestamps Arguments.of(TimestampTrait.ID, Node.from("1985-04-12T23:20:50.52Z")), Arguments.of(DateTimeTimestampTrait.ID, Node.from("1985-04-12T23:20:50.52Z")), diff --git a/smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java b/smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java index 9fa8f71a649..fc87eefd16d 100644 --- a/smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java +++ b/smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java @@ -48,6 +48,7 @@ import com.example.traits.structures.BasicAnnotationTrait; import com.example.traits.structures.NestedA; import com.example.traits.structures.NestedB; +import com.example.traits.structures.StructWithIdrefMemberTrait; import com.example.traits.structures.StructWithListOfMapTrait; import com.example.traits.structures.StructWithUniqueItemsListTrait; import com.example.traits.structures.StructureTrait; @@ -324,6 +325,13 @@ static Stream loadsModelTests() { Optional.of("a"), "getItems", Optional.of(SetUtils.of(SetUtils.of("b", "c"), SetUtils.of("d", "e"))))), + Arguments.of("structures/struct-with-idref-member-trait.smithy", + StructWithIdrefMemberTrait.class, + MapUtils.of( + "getIdRefMemberA", + Optional.of(ShapeId.from("test.smithy.traitcodegen#a")), + "getIdRefMemberB", + Optional.of(ShapeId.from("test.smithy.traitcodegen#b")))), // Timestamps Arguments.of("timestamps/struct-with-nested-timestamps.smithy", StructWithNestedTimestampsTrait.class, diff --git a/smithy-trait-codegen/src/it/resources/software/amazon/smithy/traitcodegen/test/structures/struct-with-idref-member-trait.smithy b/smithy-trait-codegen/src/it/resources/software/amazon/smithy/traitcodegen/test/structures/struct-with-idref-member-trait.smithy new file mode 100644 index 00000000000..043578e163d --- /dev/null +++ b/smithy-trait-codegen/src/it/resources/software/amazon/smithy/traitcodegen/test/structures/struct-with-idref-member-trait.smithy @@ -0,0 +1,12 @@ +$version: "2.0" + +namespace test.smithy.traitcodegen + +use test.smithy.traitcodegen.structures#StructWithIdrefMember + +@StructWithIdrefMember(idRefMemberA: "test.smithy.traitcodegen#a", idRefMemberB: "test.smithy.traitcodegen#b") +structure myStruct {} + +structure a {} + +structure b {} diff --git a/smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/BuilderGenerator.java b/smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/BuilderGenerator.java index 273b3a042f5..4f0c9d433c4 100644 --- a/smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/BuilderGenerator.java +++ b/smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/BuilderGenerator.java @@ -37,6 +37,7 @@ import software.amazon.smithy.model.shapes.UnionShape; import software.amazon.smithy.model.traits.AbstractTraitBuilder; import software.amazon.smithy.model.traits.DefaultTrait; +import software.amazon.smithy.model.traits.IdRefTrait; import software.amazon.smithy.model.traits.StringListTrait; import software.amazon.smithy.model.traits.TimestampFormatTrait; import software.amazon.smithy.model.traits.TraitDefinition; @@ -371,6 +372,10 @@ public Void mapShape(MapShape shape) { @Override public Void memberShape(MemberShape shape) { + if (shape.hasTrait(IdRefTrait.ID)) { + this.getDefault(shape); + return null; + } return model.expectShape(shape.getTarget()).accept(this); } } diff --git a/smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/FromNodeGenerator.java b/smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/FromNodeGenerator.java index 36069ca2f1e..07cea31125c 100644 --- a/smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/FromNodeGenerator.java +++ b/smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/FromNodeGenerator.java @@ -35,6 +35,7 @@ import software.amazon.smithy.model.shapes.StructureShape; import software.amazon.smithy.model.shapes.TimestampShape; import software.amazon.smithy.model.shapes.UnionShape; +import software.amazon.smithy.model.traits.IdRefTrait; import software.amazon.smithy.model.traits.TimestampFormatTrait; import software.amazon.smithy.model.traits.UniqueItemsTrait; import software.amazon.smithy.traitcodegen.SymbolProperties; @@ -242,6 +243,14 @@ private MemberGenerator(MemberShape member) { @Override public Void memberShape(MemberShape shape) { + if (shape.hasTrait(IdRefTrait.ID)) { + writer.writeInline(memberPrefix + "Member($1S, n -> $3C, builder::$2L)", + fieldName, + memberName, + (Runnable) () -> shape + .accept(new FromNodeMapperVisitor(writer, model, "n", 1, symbolProvider))); + return null; + } return model.expectShape(shape.getTarget()).accept(this); } diff --git a/smithy-trait-codegen/src/test/java/software/amazon/smithy/traitcodegen/TraitCodegenPluginTest.java b/smithy-trait-codegen/src/test/java/software/amazon/smithy/traitcodegen/TraitCodegenPluginTest.java index d0416b1d555..bce61dd23e9 100644 --- a/smithy-trait-codegen/src/test/java/software/amazon/smithy/traitcodegen/TraitCodegenPluginTest.java +++ b/smithy-trait-codegen/src/test/java/software/amazon/smithy/traitcodegen/TraitCodegenPluginTest.java @@ -26,7 +26,7 @@ import software.amazon.smithy.model.node.ObjectNode; public class TraitCodegenPluginTest { - private static final int EXPECTED_NUMBER_OF_FILES = 67; + private static final int EXPECTED_NUMBER_OF_FILES = 68; private MockManifest manifest; private Model model; diff --git a/smithy-trait-codegen/src/test/resources/META-INF/smithy/manifest b/smithy-trait-codegen/src/test/resources/META-INF/smithy/manifest index 7395407bd2a..6152a776710 100644 --- a/smithy-trait-codegen/src/test/resources/META-INF/smithy/manifest +++ b/smithy-trait-codegen/src/test/resources/META-INF/smithy/manifest @@ -42,6 +42,7 @@ structures/annotation-trait.smithy structures/structure-trait.smithy structures/struct-with-listofmap-trait.smithy structures/struct-with-uniqueitems-list-trait.smithy +structures/struct-with-idref-member-trait.smithy timestamps/date-time-format-timestamp-trait.smithy timestamps/epoch-seconds-format-timestamp-trait.smithy timestamps/http-date-format-timestamp-trait.smithy diff --git a/smithy-trait-codegen/src/test/resources/META-INF/smithy/structures/struct-with-idref-member-trait.smithy b/smithy-trait-codegen/src/test/resources/META-INF/smithy/structures/struct-with-idref-member-trait.smithy new file mode 100644 index 00000000000..b2cda7ba210 --- /dev/null +++ b/smithy-trait-codegen/src/test/resources/META-INF/smithy/structures/struct-with-idref-member-trait.smithy @@ -0,0 +1,13 @@ +$version: "2.0" + +namespace test.smithy.traitcodegen.structures + +@trait +structure StructWithIdrefMember { + @idRef(selector: "structure") + idRefMemberA: String + idRefMemberB: IdRefMemberB +} + +@idRef(selector: "structure") +string IdRefMemberB