From f8a5d1d6aea6e8ee20d092c83c1561ffccdba2f1 Mon Sep 17 00:00:00 2001
From: Michael Dowling <mtdowling@gmail.com>
Date: Thu, 6 Jan 2022 20:24:50 -0800
Subject: [PATCH] Add unwrite() to CodeWriter to help with CSVs

It's sometimes easier to deal with comma separated values by just always
throwing a ", " at the end of each element in the list, even for the
last element. The unwrite() method allows you to remove the last text
written to the CodeWriter if it matches a given substring. This
allows, for example, ", " to be "unwritten" to removing the trailing
comma without having to first check if the last text written was a
trailing comma. This is similar to Ruby's chomp() method.
---
 .../amazon/smithy/utils/CodeWriter.java       | 43 ++++++++++++++++++
 .../amazon/smithy/utils/CodeWriterTest.java   | 45 +++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/smithy-utils/src/main/java/software/amazon/smithy/utils/CodeWriter.java b/smithy-utils/src/main/java/software/amazon/smithy/utils/CodeWriter.java
index 0426853294d..6805da15e50 100644
--- a/smithy-utils/src/main/java/software/amazon/smithy/utils/CodeWriter.java
+++ b/smithy-utils/src/main/java/software/amazon/smithy/utils/CodeWriter.java
@@ -1263,6 +1263,49 @@ public final CodeWriter writeOptional(Object content) {
         }
     }
 
+    /**
+     * Remove the most recent text written to the CodeWriter if and only
+     * if the last written text is exactly equal to the given expanded
+     * content string.
+     *
+     * <p>This can be useful, for example, for use cases like removing
+     * trailing commas from lists of values.
+     *
+     * <p>For example, the following will remove ", there." from the
+     * end of the CodeWriter:
+     *
+     * <pre>{@code
+     * CodeWriter writer = new CodeWriter();
+     * writer.writeInline("Hello, there.");
+     * writer.unwrite(", there.");
+     * assert(writer.toString().equals("Hello\n"));
+     * }</pre>
+     *
+     * <p>However, the following call to unwrite will do nothing because
+     * the last text written to the CodeWriter does not match:
+     *
+     * <pre>{@code
+     * CodeWriter writer = new CodeWriter();
+     * writer.writeInline("Hello.");
+     * writer.unwrite("there.");
+     * assert(writer.toString().equals("Hello.\n"));
+     * }</pre>
+     *
+     * @param content Content to write.
+     * @param args String arguments to use for formatting.
+     * @return Returns the CodeWriter.
+     */
+    public final CodeWriter unwrite(Object content, Object... args) {
+        String value = format(content, args);
+        int currentLength = currentState.builder.length();
+
+        if (currentState.builder.lastIndexOf(value) == currentLength - value.length()) {
+            currentState.builder.setLength(currentLength - value.length());
+        }
+
+        return this;
+    }
+
     /**
      * Allows calling out to arbitrary code for things like looping or
      * conditional writes without breaking method chaining.
diff --git a/smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java b/smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java
index c25cd08b77e..ed6492013e1 100644
--- a/smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java
+++ b/smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java
@@ -776,4 +776,49 @@ public void sectionWithWriteInline() {
 
         assertThat(writer.toString(), equalTo("inline addition\n"));
     }
+
+    @Test
+    public void canUnwriteMatchingStrings() {
+        CodeWriter writer = new CodeWriter().insertTrailingNewline(false);
+        writer.writeInline("Hello there");
+        writer.unwrite(" there");
+
+        assertThat(writer.toString(), equalTo("Hello"));
+    }
+
+    @Test
+    public void unwriteDoesNothingWhenNoMatch() {
+        CodeWriter writer = new CodeWriter().insertTrailingNewline(false);
+        writer.writeInline("Hello there");
+        writer.unwrite(" nope");
+
+        assertThat(writer.toString(), equalTo("Hello there"));
+    }
+
+    @Test
+    public void canUnwriteWhenSubstringTooLong() {
+        CodeWriter writer = new CodeWriter().insertTrailingNewline(false);
+        writer.writeInline("");
+        writer.unwrite("nope");
+
+        assertThat(writer.toString(), equalTo(""));
+    }
+
+    @Test
+    public void canUnwriteWithTemplates() {
+        CodeWriter writer = new CodeWriter().insertTrailingNewline(false);
+        writer.writeInline("Hi.Hello");
+        writer.unwrite("$L", "Hello");
+
+        assertThat(writer.toString(), equalTo("Hi."));
+    }
+
+    @Test
+    public void canUnwriteWithTemplatesThatExpandToNothing() {
+        CodeWriter writer = new CodeWriter().insertTrailingNewline(false);
+        writer.writeInline("Hi.Hello");
+        writer.unwrite("$L", "");
+
+        assertThat(writer.toString(), equalTo("Hi.Hello"));
+    }
 }