Fix trait codegen to support lists of enums

smithy-lang · Oct 15, 2024 · cff969c · cff969c
1 parent 0291f02
commit cff969c
Show file tree

Hide file tree

Showing 8 changed files with 67 additions and 4 deletions.
diff --git a/...-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java b/...-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java
@@ -6,6 +6,7 @@
 import com.example.traits.defaults.StructDefaultsTrait;
 import com.example.traits.documents.DocumentTrait;
 import com.example.traits.documents.StructWithNestedDocumentTrait;
+import com.example.traits.enums.EnumListMemberTrait;
 import com.example.traits.enums.IntEnumTrait;
 import com.example.traits.enums.StringEnumTrait;
 import com.example.traits.enums.SuitTrait;
@@ -74,6 +75,9 @@ static Stream<Arguments> createTraitTests() {
                 Arguments.of(StringEnumTrait.ID, Node.from("no")),
                 Arguments.of(IntEnumTrait.ID, Node.from(2)),
                 Arguments.of(SuitTrait.ID, Node.from("clubs")),
+                Arguments.of(EnumListMemberTrait.ID, ObjectNode.objectNodeBuilder()
+                        .withMember("value", ArrayNode.fromStrings("some", "none", "some"))
+                        .build()),
                 // Lists
                 Arguments.of(NumberListTrait.ID, ArrayNode.fromNodes(
                         Node.from(1), Node.from(2), Node.from(3))

diff --git a/...rait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java b/...rait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java
@@ -7,7 +7,9 @@
 import com.example.traits.defaults.StructDefaultsTrait;
 import com.example.traits.documents.DocumentTrait;
 import com.example.traits.documents.StructWithNestedDocumentTrait;
+import com.example.traits.enums.EnumListMemberTrait;
 import com.example.traits.enums.IntEnumTrait;
+import com.example.traits.enums.SomeEnum;
 import com.example.traits.enums.StringEnumTrait;
 import com.example.traits.enums.SuitTrait;
 import com.example.traits.idref.IdRefListTrait;
@@ -95,6 +97,8 @@ static Stream<Arguments> loadsModelTests() {
                         MapUtils.of("getValue", 1, "getEnumValue", IntEnumTrait.IntEnum.YES)),
                 Arguments.of("enums/string-enum-compatibility.smithy", SuitTrait.class,
                         MapUtils.of("getEnumValue", SuitTrait.Suit.CLUB, "getValue", "club")),
+                Arguments.of("enums/enum-list-member-trait.smithy", EnumListMemberTrait.class,
+                    MapUtils.of("getValue", Optional.of(ListUtils.of(SomeEnum.SOME, SomeEnum.NONE, SomeEnum.SOME)))),
                 // Id Refs
                 Arguments.of("idref/idref-string.smithy", IdRefStringTrait.class,
                         MapUtils.of("getValue", TARGET_ONE)),

diff --git a/...it/resources/software/amazon/smithy/traitcodegen/test/enums/enum-list-member-trait.smithy b/...it/resources/software/amazon/smithy/traitcodegen/test/enums/enum-list-member-trait.smithy
@@ -0,0 +1,8 @@
+$version: "2.0"
+
+namespace test.smithy.traitcodegen
+
+use test.smithy.traitcodegen.enums#EnumListMemberTrait
+
+@EnumListMemberTrait(value: ["some", "none", "some"])
+structure myStruct {}
diff --git a/...odegen/src/main/java/software/amazon/smithy/traitcodegen/generators/BuilderGenerator.java b/...odegen/src/main/java/software/amazon/smithy/traitcodegen/generators/BuilderGenerator.java
@@ -272,22 +272,22 @@ public Void listShape(ListShape shape) {
 
             // Clear all
             writer.openBlock("public Builder clear$U() {", "}", memberName, () -> {
-                        writer.write("$L.get().clear();", memberName);
+                        writer.write("this.$L.get().clear();", memberName);
                         writer.writeWithNoFormatting("return this;");
                     }).newLine();
 
             // Set one
             writer.openBlock("public Builder add$U($T value) {", "}",
                     memberName, symbolProvider.toSymbol(shape.getMember()), () -> {
-                        writer.write("$L.get().add(value);", memberName);
+                        writer.write("this.$L.get().add(value);", memberName);
                         writer.write("return this;");
                     }).newLine();
 
             // Remove one
             writer.openBlock("public Builder remove$U($T value) {", "}",
                     memberName, symbolProvider.toSymbol(shape.getMember()),
                     () -> {
-                        writer.write("$L.get().remove(value);", memberName);
+                        writer.write("this.$L.get().remove(value);", memberName);
                         writer.write("return this;");
                     }).newLine();
             return null;

diff --git a/...degen/src/main/java/software/amazon/smithy/traitcodegen/generators/FromNodeGenerator.java b/...degen/src/main/java/software/amazon/smithy/traitcodegen/generators/FromNodeGenerator.java
@@ -36,6 +36,7 @@
 import software.amazon.smithy.model.shapes.TimestampShape;
 import software.amazon.smithy.model.shapes.UnionShape;
 import software.amazon.smithy.model.traits.TimestampFormatTrait;
+import software.amazon.smithy.traitcodegen.SymbolProperties;
 import software.amazon.smithy.traitcodegen.TraitCodegenUtils;
 import software.amazon.smithy.traitcodegen.writer.TraitCodegenWriter;
 import software.amazon.smithy.utils.StringUtils;
@@ -106,6 +107,34 @@ public Void stringShape(StringShape shape) {
         return null;
     }
 
+    @Override
+    public Void enumShape(EnumShape shape) {
+        // Enum traits do not need this method, only nested enums.
+        if (symbol.getProperty(SymbolProperties.BASE_SYMBOL).isPresent())  {
+            return null;
+        }
+        writeFromNodeJavaDoc();
+        writer.openBlock("public static $T fromNode($T node) {", "}", symbol, Node.class, () -> {
+            writer.write("return from(node.expectStringNode().getValue());");
+        });
+        writer.newLine();
+        return null;
+    }
+
+    @Override
+    public Void intEnumShape(IntEnumShape shape) {
+        // Enum traits do not need this method, only nested enums.
+        if (symbol.getProperty(SymbolProperties.BASE_SYMBOL).isPresent())  {
+            return null;
+        }
+        writeFromNodeJavaDoc();
+        writer.openBlock("public static $T fromNode($T node) {", "}", symbol, Node.class, () -> {
+            writer.writeWithNoFormatting("return from(node.expectNumberNode().getValue().intValue());");
+        });
+        writer.newLine();
+        return null;
+    }
+
     @Override
     protected Void numberShape(NumberShape shape) {
         // Number shapes do not create a from node method

diff --git a/...ait-codegen/src/test/java/software/amazon/smithy/traitcodegen/TraitCodegenPluginTest.java b/...ait-codegen/src/test/java/software/amazon/smithy/traitcodegen/TraitCodegenPluginTest.java
@@ -28,7 +28,7 @@
 
 
 public class TraitCodegenPluginTest {
-    private static final int EXPECTED_NUMBER_OF_FILES = 58;
+    private static final int EXPECTED_NUMBER_OF_FILES = 60;
 
     private MockManifest manifest;
     private Model model;

diff --git a/smithy-trait-codegen/src/test/resources/META-INF/smithy/enums/enum-list-member-trait.smithy b/smithy-trait-codegen/src/test/resources/META-INF/smithy/enums/enum-list-member-trait.smithy
@@ -0,0 +1,17 @@
+$version: "2.0"
+
+namespace test.smithy.traitcodegen.enums
+
+@trait(selector: "structure")
+structure EnumListMemberTrait {
+    value: EnumList
+}
+
+list EnumList {
+    member: SomeEnum
+}
+
+enum SomeEnum {
+    SOME = "some"
+    NONE = "none"
+}
diff --git a/smithy-trait-codegen/src/test/resources/META-INF/smithy/manifest b/smithy-trait-codegen/src/test/resources/META-INF/smithy/manifest
@@ -4,6 +4,7 @@ documents/struct-with-nested-document.smithy
 enums/enum-trait.smithy
 enums/int-enum-trait.smithy
 enums/string-enum-compatibility.smithy
+enums/enum-list-member-trait.smithy
 idref/idref-list.smithy
 idref/idref-map.smithy
 idref/idref-string.smithy