Add initial implementation of Trait initializers

codegen/src/main/java/software/amazon/smithy/java/codegen/CodeGenerationContext.java

@@ -5,13 +5,17 @@
 
 package software.amazon.smithy.java.codegen;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.codegen.core.CodegenContext;
+import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.codegen.core.WriterDelegator;
 import software.amazon.smithy.java.codegen.writer.JavaWriter;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.utils.SmithyUnstableApi;
 
 /**
@@ -27,6 +31,7 @@ public class CodeGenerationContext
     private final FileManifest fileManifest;
     private final List<JavaCodegenIntegration> integrations;
     private final WriterDelegator<JavaWriter> writerDelegator;
+    private final Map<Class<? extends Trait>, TraitInitializer> traitInitializers = new HashMap<>();
 
     CodeGenerationContext(
         Model model,
@@ -40,6 +45,7 @@ public class CodeGenerationContext
         this.symbolProvider = symbolProvider;
         this.fileManifest = fileManifest;
         this.integrations = integrations;
+        setTraitInitializers();
         this.writerDelegator = new WriterDelegator<>(fileManifest, symbolProvider, new JavaWriter.Factory(settings));
     }
 
@@ -72,4 +78,23 @@ public WriterDelegator<JavaWriter> writerDelegator() {
     public List<JavaCodegenIntegration> integrations() {
         return integrations;
     }
+
+    public TraitInitializer initializer(Class<? extends Trait> trait) {
+        return traitInitializers.get(trait);
+    }
+
+    private void setTraitInitializers() {
+        for (var integration : integrations) {
+            for (var initializer : integration.traitInitializers()) {
+                var existing = traitInitializers.put(initializer.traitClass(), initializer);
+                if (existing != null) {
+                    throw new CodegenException(
+                        "Attempted to add initializer for integration "
+                            + integration.name() + " but founding existing initializer for trait "
+                            + initializer.traitClass() + ". "
+                    );
+                }
+            }
+        }
+    }
 }

codegen/src/main/java/software/amazon/smithy/java/codegen/JavaCodegenIntegration.java

@@ -5,6 +5,8 @@
 
 package software.amazon.smithy.java.codegen;
 
+import java.util.Collections;
+import java.util.List;
 import software.amazon.smithy.codegen.core.SmithyIntegration;
 import software.amazon.smithy.java.codegen.writer.JavaWriter;
 
@@ -13,4 +15,14 @@
  * adding custom code, etc.
  */
 public interface JavaCodegenIntegration
-    extends SmithyIntegration<JavaCodegenSettings, JavaWriter, CodeGenerationContext> {}
+    extends SmithyIntegration<JavaCodegenSettings, JavaWriter, CodeGenerationContext> {
+
+    /**
+     * Gets a list of trait initializers to use when initializing traits for schemas.
+     *
+     * @return Returns the list of {@link TraitInitializer}'s.
+     */
+    default List<TraitInitializer> traitInitializers() {
+        return Collections.emptyList();
+    }
+}

codegen/src/main/java/software/amazon/smithy/java/codegen/TraitInitializer.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.codegen;
+
+import java.util.function.BiConsumer;
+import software.amazon.smithy.java.codegen.sections.TraitSection;
+import software.amazon.smithy.java.codegen.writer.JavaWriter;
+import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.traits.AnnotationTrait;
+import software.amazon.smithy.model.traits.StringListTrait;
+import software.amazon.smithy.model.traits.StringTrait;
+import software.amazon.smithy.model.traits.Trait;
+
+/**
+ * Defines how Trait classes are written in the Schema of a Shape with that Trait applied.
+ */
+public interface TraitInitializer extends BiConsumer<JavaWriter, TraitSection> {
+    Class<? extends Trait> traitClass();
+
+    record AnnotationTraitInitializer(Class<? extends AnnotationTrait> traitClass) implements TraitInitializer {
+        @Override
+        public void accept(JavaWriter writer, TraitSection section) {
+            writer.writeInline("new $T()", traitClass);
+        }
+    }
+
+    record StringTraitInitializer(Class<? extends StringTrait> traitClass) implements TraitInitializer {
+        @Override
+        public void accept(JavaWriter writer, TraitSection section) {
+            var trait = (StringTrait) section.shape().expectTrait(traitClass);
+            writer.writeInline("new $T($S)", traitClass, trait.getValue());
+        }
+    }
+
+    record StringListTraitInitializer(Class<? extends StringListTrait> traitClass) implements TraitInitializer {
+        @Override
+        public void accept(JavaWriter writer, TraitSection section) {
+            var trait = (StringListTrait) section.shape().expectTrait(traitClass);
+            writer.writeInline("new $T($S, $T.NONE)", traitClass, trait.getValues(), SourceLocation.class);
+        }
+    }
+}

codegen/src/main/java/software/amazon/smithy/java/codegen/generators/SchemaGenerator.java

@@ -7,7 +7,7 @@
 
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.java.codegen.SchemaUtils;
-import software.amazon.smithy.java.codegen.sections.SchemaTraitSection;
+import software.amazon.smithy.java.codegen.sections.TraitSection;
 import software.amazon.smithy.java.codegen.writer.JavaWriter;
 import software.amazon.smithy.java.runtime.core.schema.SdkSchema;
 import software.amazon.smithy.model.Model;
@@ -62,16 +62,15 @@ protected Void getDefault(Shape shape) {
                 """
                     static final $1T $2L = $1T.builder()
                         .type($3T.$4L)
-                        .id($5S)
-                        ${6C|}
+                        .id($5S)${6C}
                         .build();
                     """,
                 SdkSchema.class,
                 SchemaUtils.toSchemaName(shape),
                 ShapeType.class,
                 shape.getType().name(),
                 shape.toShapeId(),
-                writer.consumer(w -> writeSchemaTraitBlock(w, shape))
+                writer.consumer(w -> w.injectSection(new TraitSection(shape)))
             );
             return null;
         }
@@ -83,16 +82,15 @@ public Void listShape(ListShape shape) {
                 """
                     static final $1T $2L = $1T.builder()
                             .type($3T.LIST)
-                            .id($4S)
-                            ${5C|}
+                            .id($4S)${5C}
                             .members(SdkSchema.memberBuilder(0, "member", $6C))
                             .build();
                     """,
                 SdkSchema.class,
                 SchemaUtils.toSchemaName(shape),
                 ShapeType.class,
                 shape.toShapeId(),
-                writer.consumer(w -> writeSchemaTraitBlock(w, shape)),
+                writer.consumer(w -> w.injectSection(new TraitSection(shape))),
                 writer.consumer(w -> SchemaUtils.writeSchemaType(w, target))
             );
 
@@ -107,8 +105,7 @@ public Void mapShape(MapShape shape) {
                 """
                     static final $1T $2L = $1T.builder()
                         .type($3T.MAP)
-                        .id($4S)
-                        ${5C}
+                        .id($4S)${5C}
                         .members(
                                 $1T.memberBuilder(0, "key", $6C),
                                 $1T.memberBuilder(1, "value", $7C)
@@ -119,7 +116,7 @@ public Void mapShape(MapShape shape) {
                 SchemaUtils.toSchemaName(shape),
                 ShapeType.class,
                 shape.toShapeId(),
-                writer.consumer(w -> writeSchemaTraitBlock(w, shape)),
+                writer.consumer(w -> w.injectSection(new TraitSection(shape))),
                 writer.consumer(w -> SchemaUtils.writeSchemaType(w, keyShape)),
                 writer.consumer(w -> SchemaUtils.writeSchemaType(w, valueShape))
             );
@@ -146,8 +143,7 @@ public Void structureShape(StructureShape shape) {
                 """
                     static final $1T SCHEMA = $1T.builder()
                         .id(ID)
-                        .type($2T.$3L)
-                        ${4C|}
+                        .type($2T.$3L)${4C}
                         ${?memberSchemas}.members(${#memberSchemas}
                             ${value:L}${^key.last},${/key.last}${/memberSchemas}
                         )
@@ -156,7 +152,7 @@ public Void structureShape(StructureShape shape) {
                 SdkSchema.class,
                 ShapeType.class,
                 shape.getType().toString().toUpperCase(),
-                writer.consumer(w -> writeSchemaTraitBlock(w, shape))
+                writer.consumer(w -> w.injectSection(new TraitSection(shape)))
             );
             writer.popState();
 
@@ -174,27 +170,15 @@ private void writeNestedMemberSchema(int idx, MemberShape member) {
             writer.write(
                 """
                     private static final $1T $2L = $1T.memberBuilder($3L, $4S, $5C)
-                        .id(ID)
-                        ${6C|}
+                        .id(ID)${6C}
                         .build();
                     """,
                 SdkSchema.class,
                 SchemaUtils.toMemberSchemaName(memberName),
                 idx,
                 memberName,
                 writer.consumer(w -> SchemaUtils.writeSchemaType(w, target)),
-                writer.consumer(w -> writeSchemaTraitBlock(w, member))
-            );
-        }
-
-        private static void writeSchemaTraitBlock(JavaWriter writer, Shape shape) {
-            if (shape.getAllTraits().isEmpty()) {
-                return;
-            }
-            writer.openBlock(
-                ".traits(",
-                ")",
-                () -> writer.injectSection(new SchemaTraitSection(shape)).newLine()
+                writer.consumer(w -> w.injectSection(new TraitSection(member)))
             );
         }
     }

codegen/src/main/java/software/amazon/smithy/java/codegen/integrations/core/CoreIntegration.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.codegen.integrations.core;
+
+import java.util.List;
+import software.amazon.smithy.java.codegen.CodeGenerationContext;
+import software.amazon.smithy.java.codegen.JavaCodegenIntegration;
+import software.amazon.smithy.java.codegen.TraitInitializer;
+import software.amazon.smithy.java.codegen.writer.JavaWriter;
+import software.amazon.smithy.model.traits.ErrorTrait;
+import software.amazon.smithy.model.traits.HttpHeaderTrait;
+import software.amazon.smithy.model.traits.HttpQueryParamsTrait;
+import software.amazon.smithy.model.traits.HttpResponseCodeTrait;
+import software.amazon.smithy.model.traits.IdempotentTrait;
+import software.amazon.smithy.model.traits.InputTrait;
+import software.amazon.smithy.model.traits.JsonNameTrait;
+import software.amazon.smithy.model.traits.OutputTrait;
+import software.amazon.smithy.model.traits.ReadonlyTrait;
+import software.amazon.smithy.model.traits.RequiredTrait;
+import software.amazon.smithy.model.traits.SensitiveTrait;
+import software.amazon.smithy.model.traits.StreamingTrait;
+import software.amazon.smithy.model.traits.TagsTrait;
+import software.amazon.smithy.model.traits.TimestampFormatTrait;
+import software.amazon.smithy.utils.CodeInterceptor;
+import software.amazon.smithy.utils.CodeSection;
+
+/**
+ * Base Plugin for
+ * TODO: more complete docs once this has all the expected interceptors
+ */
+public class CoreIntegration implements JavaCodegenIntegration {
+
+    @Override
+    public String name() {
+        return "core";
+    }
+
+    @Override
+    public List<? extends CodeInterceptor<? extends CodeSection, JavaWriter>> interceptors(
+        CodeGenerationContext context
+    ) {
+        return List.of(
+            new TraitInitializerInterceptor(context)
+        );
+    }
+
+    @Override
+    public List<TraitInitializer> traitInitializers() {
+        return List.of(
+            new TraitInitializer.AnnotationTraitInitializer(RequiredTrait.class),
+            new TraitInitializer.AnnotationTraitInitializer(SensitiveTrait.class),
+            new TraitInitializer.StringTraitInitializer(JsonNameTrait.class),
+            new TraitInitializer.AnnotationTraitInitializer(StreamingTrait.class),
+            new TraitInitializer.AnnotationTraitInitializer(HttpResponseCodeTrait.class),
+            new TraitInitializer.StringTraitInitializer(ErrorTrait.class),
+            new TraitInitializer.AnnotationTraitInitializer(HttpQueryParamsTrait.class),
+            new TraitInitializer.StringTraitInitializer(TimestampFormatTrait.class),
+            new TraitInitializer.AnnotationTraitInitializer(IdempotentTrait.class),
+            new TraitInitializer.AnnotationTraitInitializer(ReadonlyTrait.class),
+            new TraitInitializer.StringListTraitInitializer(TagsTrait.class),
+            new TraitInitializer.StringTraitInitializer(HttpHeaderTrait.class),
+            new TraitInitializer.AnnotationTraitInitializer(InputTrait.class),
+            new TraitInitializer.AnnotationTraitInitializer(OutputTrait.class)
+        );
+    }
+}

codegen/src/main/java/software/amazon/smithy/java/codegen/integrations/core/TraitInitializerInterceptor.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.codegen.integrations.core;
+
+import java.util.HashSet;
+import java.util.Set;
+import software.amazon.smithy.java.codegen.CodeGenerationContext;
+import software.amazon.smithy.java.codegen.TraitInitializer;
+import software.amazon.smithy.java.codegen.sections.TraitSection;
+import software.amazon.smithy.java.codegen.writer.JavaWriter;
+import software.amazon.smithy.utils.CodeInterceptor;
+
+/**
+ * Injects Trait definitions using trait initializers.
+ *
+ * @param context Code generation context.
+ */
+record TraitInitializerInterceptor(CodeGenerationContext context) implements CodeInterceptor<TraitSection, JavaWriter> {
+    private static final System.Logger LOGGER = System.getLogger(TraitInitializerInterceptor.class.getName());
+
+    @Override
+    public Class<TraitSection> sectionType() {
+        return TraitSection.class;
+    }
+
+    @Override
+    public void write(JavaWriter writer, String previousText, TraitSection section) {
+        if (section.shape().getAllTraits().isEmpty()) {
+            // Do not print anything if there are no traits on the shape.
+            return;
+        }
+
+        // Collect all initializers
+        Set<TraitInitializer> initializers = new HashSet<>();
+        for (var trait : section.shape().getAllTraits().values()) {
+            var initializer = context.initializer(trait.getClass());
+            if (initializer == null) {
+                LOGGER.log(System.Logger.Level.DEBUG, "No trait initializer found for %s. Skipping...", trait);
+                continue;
+            }
+            initializers.add(initializer);
+        }
+
+        // Do not print anything if there are no initializers
+        if (initializers.isEmpty()) {
+            return;
+        }
+
+        // Add the method call for traits now we know valid initializers exist
+        writer.newLine();
+        writer.indent();
+        writer.openBlock(".traits(", ")", () -> {
+            // Loop through initializers to inject trait definitions
+            var iter = initializers.iterator();
+            while (iter.hasNext()) {
+                iter.next().accept(writer, section);
+                if (iter.hasNext()) {
+                    writer.writeInlineWithNoFormatting(", \n");
+                }
+            }
+            writer.newLine();
+        });
+        writer.dedent();
+    }
+}

codegen/src/main/java/software/amazon/smithy/java/codegen/sections/SchemaTraitSection.java → codegen/src/main/java/software/amazon/smithy/java/codegen/sections/TraitSection.java

@@ -7,10 +7,12 @@
 
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.utils.CodeSection;
+import software.amazon.smithy.utils.SmithyInternalApi;
 
 /**
- * Injected section used to populate trait values for SdkSchemas
+ * Injects Trait definitions for a Smithy Shape.
  *
- * @param shape Shape traits are a attached to.
+ * @param shape Smithy shape to inject trait definitions for.
  */
-public record SchemaTraitSection(Shape shape) implements CodeSection {}
+@SmithyInternalApi
+public record TraitSection(Shape shape) implements CodeSection {}