From 03cd10624d750a5fc04eede5b7f6c7143ab61abb Mon Sep 17 00:00:00 2001
From: Michael Dowling <mtdowling@gmail.com>
Date: Wed, 30 Mar 2022 22:01:49 -0700
Subject: [PATCH] Add ability to register interceptors w/ delegator

---
 .../smithy/codegen/core/WriterDelegator.java  | 19 +++++++++++++-
 .../codegen/core/WriterDelegatorTest.java     | 26 +++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/WriterDelegator.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/WriterDelegator.java
index 2fc3f6ac474..0156be49ce4 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/WriterDelegator.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/WriterDelegator.java
@@ -25,6 +25,8 @@
 import java.util.function.Consumer;
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.utils.CodeInterceptor;
+import software.amazon.smithy.utils.CodeSection;
 
 /**
  * <p>Creates and manages {@link SymbolWriter}s for files and namespaces based
@@ -66,6 +68,7 @@ public class WriterDelegator<W extends SymbolWriter<W, ? extends ImportContainer
     private final SymbolProvider symbolProvider;
     private final Map<String, W> writers = new TreeMap<>();
     private final SymbolWriter.Factory<W> factory;
+    private final List<CodeInterceptor<? extends CodeSection, W>> interceptors = new ArrayList<>();
     private String automaticSeparator = "\n";
 
     /**
@@ -79,6 +82,16 @@ public WriterDelegator(FileManifest fileManifest, SymbolProvider symbolProvider,
         this.factory = factory;
     }
 
+    /**
+     * Sets the list of code interceptors to registered with each newly created writer.
+     *
+     * @param interceptors Interceptors to register.
+     */
+    public void setInterceptors(List<? extends CodeInterceptor<? extends CodeSection, W>> interceptors) {
+        this.interceptors.clear();
+        this.interceptors.addAll(interceptors);
+    }
+
     /**
      * Gets all of the dependencies that have been registered in writers
      * created by the {@code WriterDelegator}.
@@ -214,7 +227,11 @@ private W checkoutWriter(String filename, String namespace) {
         String formattedFilename = Paths.get(filename).normalize().toString();
         boolean needsNewline = writers.containsKey(formattedFilename);
 
-        W writer = writers.computeIfAbsent(formattedFilename, file -> factory.apply(file, namespace));
+        W writer = writers.computeIfAbsent(formattedFilename, file -> {
+            W result = factory.apply(file, namespace);
+            interceptors.forEach(result::onSection);
+            return result;
+        });
 
         // Add newlines/separators between types in the same file.
         if (needsNewline) {
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/WriterDelegatorTest.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/WriterDelegatorTest.java
index d287bcd1d08..4a5f48de078 100644
--- a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/WriterDelegatorTest.java
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/WriterDelegatorTest.java
@@ -26,6 +26,9 @@
 import software.amazon.smithy.build.MockManifest;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.StringShape;
+import software.amazon.smithy.utils.CodeInterceptor;
+import software.amazon.smithy.utils.CodeSection;
+import software.amazon.smithy.utils.ListUtils;
 
 public class WriterDelegatorTest {
     @Test
@@ -120,4 +123,27 @@ public void flushesAllWriters() {
 
         assertThat(mockManifest.getFileString("com/foo/Baz.bam"), equalTo(Optional.of("Hi!\n")));
     }
+
+    @Test
+    public void registersInterceptors() {
+        MockManifest mockManifest = new MockManifest();
+        SymbolProvider provider = (shape) -> Symbol.builder()
+                .namespace("com.foo", ".")
+                .name("Baz")
+                .definitionFile("com/foo/Baz.bam")
+                .build();
+        WriterDelegator<MySimpleWriter> delegator = new WriterDelegator<>(
+                mockManifest, provider, (f, n) -> new MySimpleWriter(n));
+
+        CodeInterceptor<CodeSection, MySimpleWriter> a = CodeInterceptor.forName("test", (w, s) -> {
+            w.write("Yes");
+        });
+        delegator.setInterceptors(ListUtils.of(a));
+
+        delegator.useFileWriter("/foo.txt", w -> {
+            w.pushState("test").popState();
+        });
+
+        assertThat(delegator.getWriters().get(Paths.get("/foo.txt").toString()).toString(), equalTo("Yes\n"));
+    }
 }