Refactor writeAdditionalFiles and writeAdditionalExports logic into integration.customize()

AndrewFossAWS · AndrewFossAWS · commit 88b4e8e74efb · 2022-10-11T10:23:46.000-07:00
diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/DirectedTypeScriptCodegen.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/DirectedTypeScriptCodegen.java
@@ -22,6 +22,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import software.amazon.smithy.build.FileManifest;
@@ -404,24 +406,19 @@ public void customizeBeforeIntegrations(
             }
         }
 
-        // Write each custom file.
-        for (TypeScriptIntegration integration : directive.context().integrations()) {
-            LOGGER.finer(() -> "Calling writeAdditionalFiles on " + integration.getClass().getCanonicalName());
-            integration.writeAdditionalFiles(
-                    directive.settings(),
-                    directive.model(),
-                    directive.symbolProvider(),
-                    directive.context().writerDelegator()::useFileWriter);
-        }
-
         // Generate index for client.
-        IndexGenerator.writeIndex(
+        BiConsumer<String, Consumer<TypeScriptWriter>> writerFactory =
+            directive.context().writerDelegator()::useFileWriter;
+
+        writerFactory.accept(Paths.get(CodegenUtils.SOURCE_FOLDER, "index.ts").toString(), writer -> {
+            IndexGenerator.writeIndex(
                 directive.settings(),
                 directive.model(),
                 directive.symbolProvider(),
-                directive.fileManifest(),
-                directive.context().integrations(),
-                directive.context().protocolGenerator());
+                directive.context().protocolGenerator(),
+                writer
+            );
+        });
 
         if (directive.settings().generateServerSdk()) {
             // Generate index for server
diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/IndexGenerator.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/IndexGenerator.java
@@ -27,7 +27,6 @@
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.traits.PaginatedTrait;
 import software.amazon.smithy.typescript.codegen.integration.ProtocolGenerator;
-import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
 import software.amazon.smithy.utils.SmithyInternalApi;
 import software.amazon.smithy.waiters.WaitableTrait;
 
@@ -43,14 +42,11 @@ static void writeIndex(
         TypeScriptSettings settings,
         Model model,
         SymbolProvider symbolProvider,
-        FileManifest fileManifest,
-        List<TypeScriptIntegration> integrations,
-        ProtocolGenerator protocolGenerator
+        ProtocolGenerator protocolGenerator,
+        TypeScriptWriter writer
     ) {
-        TypeScriptWriter writer = new TypeScriptWriter("");
-
         if (settings.generateClient()) {
-            writeClientExports(settings, model, symbolProvider, writer, fileManifest, integrations);
+            writeClientExports(settings, model, symbolProvider, writer);
         }
 
         if (settings.generateServerSdk() && protocolGenerator != null) {
@@ -60,13 +56,6 @@ static void writeIndex(
 
         // write export statement for models
         writer.write("export * from \"./models\";");
-
-        // Write each custom export.
-        for (TypeScriptIntegration integration : integrations) {
-            integration.writeAdditionalExports(settings, model, symbolProvider, writer);
-        }
-
-        fileManifest.writeFile(Paths.get(CodegenUtils.SOURCE_FOLDER, "index.ts").toString(), writer.toString());
     }
 
     private static void writeProtocolExports(ProtocolGenerator protocolGenerator, TypeScriptWriter writer) {
@@ -97,9 +86,7 @@ private static void writeClientExports(
             TypeScriptSettings settings,
             Model model,
             SymbolProvider symbolProvider,
-            TypeScriptWriter writer,
-            FileManifest fileManifest,
-            List<TypeScriptIntegration> integrations
+            TypeScriptWriter writer
     ) {
         ServiceShape service = settings.getService(model);
         Symbol symbol = symbolProvider.toSymbol(service);
diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddBaseServiceExceptionClass.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddBaseServiceExceptionClass.java
@@ -25,6 +25,7 @@
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.traits.ErrorTrait;
 import software.amazon.smithy.typescript.codegen.CodegenUtils;
+import software.amazon.smithy.typescript.codegen.TypeScriptCodegenContext;
 import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
 import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
 import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
@@ -37,7 +38,20 @@
 public final class AddBaseServiceExceptionClass implements TypeScriptIntegration {
 
     @Override
-    public void writeAdditionalFiles(
+    public void customize(TypeScriptCodegenContext codegenContext) {
+        TypeScriptSettings settings = codegenContext.settings();
+        Model model = codegenContext.model();
+        SymbolProvider symbolProvider = codegenContext.symbolProvider();
+        BiConsumer<String, Consumer<TypeScriptWriter>> writerFactory = codegenContext.writerDelegator()::useFileWriter;
+
+        writeAdditionalFiles(settings, model, symbolProvider, writerFactory);
+
+        writerFactory.accept(Paths.get(CodegenUtils.SOURCE_FOLDER, "index.ts").toString(), writer -> {
+            writeAdditionalExports(settings, model, symbolProvider, writer);
+        });
+    }
+
+    private void writeAdditionalFiles(
             TypeScriptSettings settings,
             Model model,
             SymbolProvider symbolProvider,
@@ -67,8 +81,7 @@ public void writeAdditionalFiles(
         }
     }
 
-    @Override
-    public void writeAdditionalExports(
+    private void writeAdditionalExports(
             TypeScriptSettings settings,
             Model model,
             SymbolProvider symbolProvider,
diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddHttpApiKeyAuthPlugin.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddHttpApiKeyAuthPlugin.java
@@ -23,7 +23,6 @@
 import java.util.function.Consumer;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolDependency;
-import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
 import software.amazon.smithy.model.knowledge.TopDownIndex;
@@ -32,6 +31,7 @@
 import software.amazon.smithy.model.traits.HttpApiKeyAuthTrait;
 import software.amazon.smithy.model.traits.OptionalAuthTrait;
 import software.amazon.smithy.typescript.codegen.CodegenUtils;
+import software.amazon.smithy.typescript.codegen.TypeScriptCodegenContext;
 import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
 import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
 import software.amazon.smithy.utils.IoUtils;
@@ -101,10 +101,21 @@ public List<RuntimeClientPlugin> getClientPlugins() {
     }
 
     @Override
-    public void writeAdditionalFiles(
+    public void customize(TypeScriptCodegenContext codegenContext) {
+        TypeScriptSettings settings = codegenContext.settings();
+        Model model = codegenContext.model();
+        BiConsumer<String, Consumer<TypeScriptWriter>> writerFactory = codegenContext.writerDelegator()::useFileWriter;
+
+        writeAdditionalFiles(settings, model, writerFactory);
+
+        writerFactory.accept(Paths.get(CodegenUtils.SOURCE_FOLDER, "index.ts").toString(), writer -> {
+            writeAdditionalExports(settings, model, writer);
+        });
+    }
+
+    private void writeAdditionalFiles(
         TypeScriptSettings settings,
         Model model,
-        SymbolProvider symbolProvider,
         BiConsumer<String, Consumer<TypeScriptWriter>> writerFactory
     ) {
         ServiceShape service = settings.getService(model);
@@ -144,11 +155,9 @@ public void writeAdditionalFiles(
                 });
     }
 
-    @Override
-    public void writeAdditionalExports(
+    private void writeAdditionalExports(
             TypeScriptSettings settings,
             Model model,
-            SymbolProvider symbolProvider,
             TypeScriptWriter writer
     ) {
         boolean isClientSdk = settings.generateClient();
diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/TypeScriptIntegration.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/TypeScriptIntegration.java
@@ -18,7 +18,6 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import software.amazon.smithy.codegen.core.SmithyIntegration;
 import software.amazon.smithy.codegen.core.Symbol;
@@ -41,75 +40,6 @@
 public interface TypeScriptIntegration
         extends SmithyIntegration<TypeScriptSettings, TypeScriptWriter, TypeScriptCodegenContext> {
 
-    /**
-     * Writes additional files.
-     *
-     * <pre>
-     * {@code
-     * public final class MyIntegration implements TypeScriptIntegration {
-     *     public writeAdditionalFiles(
-     *             TypeScriptSettings settings,
-     *             Model model,
-     *             SymbolProvider symbolProvider,
-     *             BiConsumer<String, Consumer<TypeScriptWriter>> writerFactory
-     *     ) {
-     *         writerFactory.accept("foo.ts", writer -> {
-     *             writer.write("// Hello!");
-     *         });
-     *     }
-     * }
-     * }</pre>
-     *
-     * @param settings Settings used to generate.
-     * @param model Model to generate from.
-     * @param symbolProvider Symbol provider used for codegen.
-     * @param writerFactory A factory function that takes the name of a file
-     *   to write and a {@code Consumer} that receives a
-     *   {@link TypeScriptWriter} to perform the actual writing to the file.
-     */
-    default void writeAdditionalFiles(
-            TypeScriptSettings settings,
-            Model model,
-            SymbolProvider symbolProvider,
-            BiConsumer<String, Consumer<TypeScriptWriter>> writerFactory
-    ) {
-        // pass
-    }
-
-    /**
-     * Write additional export statement to client root index.ts.
-     * The files generated by {@link TypeScriptIntegration#writeAdditionalFiles}
-     * can be exported with this method.
-     *
-     * <pre>
-     * {@code
-     * public final class MyIntegration implements TypeScriptIntegration {
-     *     public writeAdditionalExports(
-     *             TypeScriptSettings settings,
-     *             Model model,
-     *             SymbolProvider symbolProvider,
-     *             TypeScriptWriter writer
-     *     ) {
-     *         writer.write("export * from $S;", "./foo.ts");
-     *     }
-     * }
-     * }</pre>
-     *
-     * @param settings Settings used to generate.
-     * @param model Model to generate from.
-     * @param symbolProvider Symbol provider used for codegen.
-     * @param writer TypeScript writer that works with the file.
-     * @see #writeAdditionalFiles(TypeScriptSettings, Model, SymbolProvider, BiConsumer)
-     */
-    default void writeAdditionalExports(
-            TypeScriptSettings settings,
-            Model model,
-            SymbolProvider symbolProvider,
-            TypeScriptWriter writer
-    ) {
-        // pass
-    }
-
     /**
      * Gets a list of plugins to apply to the generated client.
      *
diff --git a/smithy-typescript-codegen/src/test/java/software/amazon/smithy/typescript/codegen/IndexGeneratorTest.java b/smithy-typescript-codegen/src/test/java/software/amazon/smithy/typescript/codegen/IndexGeneratorTest.java
@@ -3,14 +3,10 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
 
-import java.util.ArrayList;
-import java.util.List;
 import org.junit.jupiter.api.Test;
-import software.amazon.smithy.build.MockManifest;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.Node;
-import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
 
 public class IndexGeneratorTest {
 
@@ -23,27 +19,14 @@ public void writesIndex() {
                 .withMember("packageVersion", Node.from("1.0.0"))
                 .build());
         SymbolProvider symbolProvider = new SymbolVisitor(model, settings);
-        MockManifest manifest = new MockManifest();
-        List<TypeScriptIntegration> integrations = new ArrayList<>();
-        integrations.add(new TypeScriptIntegration() {
-            @Override
-            public void writeAdditionalExports(
-                    TypeScriptSettings settings,
-                    Model model,
-                    SymbolProvider symbolProvider,
-                    TypeScriptWriter writer
-            ) {
-                writer.write("export * from $S;", "./foo");
-            }
-        });
+        TypeScriptWriter writer = new TypeScriptWriter("");
 
-        IndexGenerator.writeIndex(settings, model, symbolProvider, manifest, integrations, null);
+        IndexGenerator.writeIndex(settings, model, symbolProvider, null, writer);
 
-        String contents = manifest.getFileString(CodegenUtils.SOURCE_FOLDER + "/index.ts").get();
+        String contents = writer.toString();
         assertThat(contents, containsString("export * from \"./Example\";"));
         assertThat(contents, containsString("export * from \"./ExampleClient\";"));
         assertThat(contents, containsString("export * from \"./commands\";"));
         assertThat(contents, containsString("export * from \"./models\";"));
-        assertThat(contents, containsString("export * from \"./foo\";"));
     }
 }
