From 05179138a21d6e0c9795f09f5a6bd466edac7398 Mon Sep 17 00:00:00 2001
From: Joshua Barr <jobarr@amazon.com>
Date: Fri, 1 Nov 2024 13:53:46 -0700
Subject: [PATCH] Simplify OutputFormat

---
 .../com/amazon/tools/cli/OutputFormat.java    | 114 +++++++-----------
 1 file changed, 41 insertions(+), 73 deletions(-)

diff --git a/ion-java-cli/src/main/java/com/amazon/tools/cli/OutputFormat.java b/ion-java-cli/src/main/java/com/amazon/tools/cli/OutputFormat.java
index 5ccc16d2e..948e55d25 100644
--- a/ion-java-cli/src/main/java/com/amazon/tools/cli/OutputFormat.java
+++ b/ion-java-cli/src/main/java/com/amazon/tools/cli/OutputFormat.java
@@ -4,85 +4,53 @@
 import com.amazon.ion.SymbolTable;
 import com.amazon.ion.system.IonBinaryWriterBuilder;
 import com.amazon.ion.system.IonTextWriterBuilder;
+import com.amazon.ion.system.IonWriterBuilder;
 
 import java.io.OutputStream;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 /**
  * Represents the different Ion output formats supported by the command line com.amazon.tools in this package.
  */
 public enum OutputFormat {
-    /**
-     * Nicely spaced, 'prettified' text Ion.
-     */
-    PRETTY {
-        @Override
-        public IonWriter createIonWriter(OutputStream outputStream) {
-            return IonTextWriterBuilder.pretty().build(outputStream);
-        }
-
-        @Override
-        public IonWriter createIonWriterWithImports(OutputStream outputStream, SymbolTable[] imports) {
-            return IonTextWriterBuilder.pretty().withImports(imports).build(outputStream);
-        }
-    },
-    /**
-     * Minimally spaced text Ion.
-     */
-    TEXT {
-        @Override
-        public IonWriter createIonWriter(OutputStream outputStream) {
-            return IonTextWriterBuilder.standard().build(outputStream);
-        }
-
-        @Override
-        public IonWriter createIonWriterWithImports(OutputStream outputStream, SymbolTable[] imports) {
-            return IonTextWriterBuilder.standard().withImports(imports).build(outputStream);
-        }
-    },
-    /**
-     * Compact, read-optimized binary Ion.
-     */
-    BINARY {
-        @Override
-        public IonWriter createIonWriter(OutputStream outputStream) {
-            return IonBinaryWriterBuilder.standard().build(outputStream);
-        }
-
-        @Override
-        public IonWriter createIonWriterWithImports(OutputStream outputStream, SymbolTable[] imports) {
-            return IonBinaryWriterBuilder.standard().withImports(imports).build(outputStream);
-        }
-    },
-    /**
-     * Event Stream
-     */
-    EVENTS {
-        @Override
-        public IonWriter createIonWriter(OutputStream outputStream) {
-            return IonTextWriterBuilder.pretty().build(outputStream);
-        }
-
-        @Override
-        public IonWriter createIonWriterWithImports(OutputStream outputStream, SymbolTable[] imports) {
-            return IonTextWriterBuilder.pretty().withImports(imports).build(outputStream);
-        }
-    },
-    /**
-     * None
-     */
-    NONE {
-        @Override
-        public IonWriter createIonWriter(OutputStream outputStream) {
-            NoOpOutputStream out = new NoOpOutputStream();
-            return IonTextWriterBuilder.standard().build(out);
-        }
-
-        @Override
-        public IonWriter createIonWriterWithImports(OutputStream outputStream, SymbolTable[] imports) {
-            return IonTextWriterBuilder.standard().withImports(imports).build(outputStream);
-        }
-    };
+    /** Nicely spaced, 'prettified' text Ion */
+    PRETTY(IonTextWriterBuilder::pretty),
+    /** Minimally spaced text Ion */
+    TEXT(IonTextWriterBuilder::standard),
+    /** Compact, read-optimized binary Ion */
+    BINARY(IonBinaryWriterBuilder::standard),
+    /** Event Stream */
+    EVENTS(IonTextWriterBuilder::pretty),
+    /** No output, /dev/null */
+    NONE(IonTextWriterBuilder::standard, o -> new NoOpOutputStream());
+
+    OutputFormat(Supplier<IonWriterBuilder> supplier) {
+        this(supplier, o -> o);
+    }
+
+    OutputFormat(Supplier<IonWriterBuilder> supplier, Function<OutputStream, OutputStream> function) {
+        this.writerBuilderSupplier = supplier;
+        this.outputTransformer = function;
+    }
+
+    IonWriter createIonWriter(OutputStream outputStream) {
+        return writerBuilderSupplier.get().build(outputTransformer.apply(outputStream));
+    }
+
+    IonWriter createIonWriterWithImports(OutputStream outputStream, SymbolTable[] symbolTable) {
+        IonWriterBuilder builder = writerBuilderSupplier.get();
+        OutputStream out = outputTransformer.apply(outputStream);
+        if (builder instanceof IonTextWriterBuilder) {
+            return ((IonTextWriterBuilder)builder).withImports(symbolTable).build(out);
+        } else if (builder instanceof IonBinaryWriterBuilder) {
+            return ((IonBinaryWriterBuilder)builder).withImports(symbolTable).build(out);
+        } else {
+            throw new IllegalStateException("This is impossible, none of the enums can do this");
+        }
+    }
+
+    private final Supplier<IonWriterBuilder> writerBuilderSupplier;
+    private final Function<OutputStream, OutputStream> outputTransformer;
 
-    abstract IonWriter createIonWriter(OutputStream outputStream);
-    abstract IonWriter createIonWriterWithImports(OutputStream outputStream, SymbolTable[] symbolTable);
 }