Add sourceLocation setting in fromNode() in trait codegen

Author: joewyz

smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java

@@ -71,6 +71,7 @@
 public class CreatesTraitTest {
     private static final ShapeId DUMMY_ID = ShapeId.from("ns.foo#foo");
     private final TraitFactory provider = TraitFactory.createServiceFactory();
+    private static final SourceLocation testLocation = new SourceLocation("test.smithy", 1, 2);
 
     static Stream<Arguments> createTraitTests() {
         return Stream.of(
@@ -236,4 +237,44 @@ void createsTraitFromNode(ShapeId traitId, Node fromNode) {
         assertEquals(SourceLocation.NONE, trait.getSourceLocation());
         assertEquals(trait, provider.createTrait(traitId, DUMMY_ID, trait.toNode()).orElseThrow(RuntimeException::new));
     }
+
+    static Stream<Arguments> createSourceLocationTests() {
+        return Stream.of(
+                Arguments.of(NumberListTrait.ID,
+                        ArrayNode.builder()
+                                .withValue(Node.from(1))
+                                .withValue(Node.from(2))
+                                .withValue(Node.from(3))
+                                .sourceLocation(testLocation)
+                                .build()
+                                .toNode()),
+                Arguments.of(NumberSetTrait.ID,
+                        ArrayNode.builder()
+                                .withValue(Node.from(1))
+                                .withValue(Node.from(2))
+                                .withValue(Node.from(3))
+                                .sourceLocation(testLocation)
+                                .build()
+                                .toNode()),
+                Arguments.of(NestedMapTrait.ID,
+                        NestedMapTrait.builder()
+                                .putValues("1", MapUtils.of("1", MapUtils.of("2", "3")))
+                                .sourceLocation(testLocation)
+                                .build()
+                                .toNode()),
+                Arguments.of(StructWithListOfMapTrait.ID,
+                        StructWithListOfMapTrait.builder()
+                                .addItems(MapUtils.of("1", "2"))
+                                .addItems(MapUtils.of("3", "4"))
+                                .sourceLocation(testLocation)
+                                .build()
+                                .toNode()));
+    }
+
+    @ParameterizedTest
+    @MethodSource("createSourceLocationTests")
+    void sourceLocationTest(ShapeId traitId, Node fromNode) {
+        Trait trait = provider.createTrait(traitId, DUMMY_ID, fromNode).orElseThrow(RuntimeException::new);
+        assertEquals(testLocation, trait.getSourceLocation());
+    }
 }

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/FromNodeGenerator.java

@@ -37,6 +37,7 @@
 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.TraitDefinition;
 import software.amazon.smithy.model.traits.UniqueItemsTrait;
 import software.amazon.smithy.traitcodegen.SymbolProperties;
 import software.amazon.smithy.traitcodegen.TraitCodegenUtils;
@@ -83,7 +84,8 @@ public Void listShape(ListShape shape) {
                 symbol,
                 Node.class,
                 () -> {
-                    writer.writeWithNoFormatting("Builder builder = builder();");
+                    writer.writeWithNoFormatting("Builder builder = builder()" +
+                            (shape.hasTrait(TraitDefinition.ID) ? ".sourceLocation(node);" : ";"));
                     shape.accept(new FromNodeMapperVisitor(writer, model, "node", symbolProvider));
                     writer.write("builder.values(value0);");
                     writer.writeWithNoFormatting("return builder.build();");
@@ -101,7 +103,8 @@ public Void mapShape(MapShape shape) {
                 symbol,
                 Node.class,
                 () -> {
-                    writer.writeWithNoFormatting("Builder builder = builder();");
+                    writer.writeWithNoFormatting("Builder builder = builder()" +
+                            (shape.hasTrait(TraitDefinition.ID) ? ".sourceLocation(node);" : ";"));
                     shape.accept(new FromNodeMapperVisitor(writer, model, "node", symbolProvider));
                     writer.writeWithNoFormatting("return builder.build();");
                 });
@@ -159,7 +162,8 @@ public Void structureShape(StructureShape shape) {
         writeFromNodeJavaDoc();
         writer.write("public static $T fromNode($T node) {", symbol, Node.class);
         writer.indent();
-        writer.write("Builder builder = builder();");
+        writer.write("Builder builder = builder()" +
+                (shape.hasTrait(TraitDefinition.ID) ? ".sourceLocation(node);" : ";"));
         // If the shape has no members (i.e. is an annotation trait) then there will be no member setters, and we
         // need to terminate the line.
         writer.putContext("isEmpty", shape.members().isEmpty());