Update OutputFormatter

rundeck · Nov 18, 2016 · 6f073c1 · 6f073c1
1 parent ef427dc
commit 6f073c1
Show file tree

Hide file tree

Showing 6 changed files with 55 additions and 6 deletions.
diff --git a/...snakeyaml/src/main/java/com/simplifyops/toolbelt/format/yaml/snakeyaml/YamlFormatter.java b/...snakeyaml/src/main/java/com/simplifyops/toolbelt/format/yaml/snakeyaml/YamlFormatter.java
@@ -57,6 +57,10 @@ private YamlFormatter(Yaml yaml, final OutputFormatter base) {
         this.base = base;
     }
 
+    @Override
+    public OutputFormatter withBase(final OutputFormatter base) {
+        return new YamlFormatter(this.yaml, base);
+    }
 
     /**
      * @param base    base formatter

diff --git a/toolbelt/src/main/java/com/simplifyops/toolbelt/ANSIColorOutput.java b/toolbelt/src/main/java/com/simplifyops/toolbelt/ANSIColorOutput.java
@@ -11,12 +11,23 @@ public class ANSIColorOutput implements CommandOutput, OutputFormatter {
     private static final Object RED = ESC + "[31m";
     private static final Object YELLOW = ESC + "[33m";
 
+    private OutputFormatter base;
     SystemOutput sink;
 
     public ANSIColorOutput(final SystemOutput sink) {
         this.sink = sink;
     }
 
+    public ANSIColorOutput(final SystemOutput sink, OutputFormatter base) {
+        this.sink = sink;
+        this.base = base;
+    }
+
+    @Override
+    public OutputFormatter withBase(final OutputFormatter base) {
+        return new ANSIColorOutput(sink, base);
+    }
+
     @Override
     public String format(final Object o) {
         return toColors(o);
@@ -28,13 +39,17 @@ public void output(final Object object) {
     }
 
     public static String toColors(final Object object) {
+        return toColors(object, null);
+    }
+
+    public static String toColors(final Object object, OutputFormatter base) {
         if (null == object) {
             return null;
         }
         if (ColorString.class.isAssignableFrom(object.getClass())) {
             ColorString object1 = (ColorString) object;
             Set<ColorArea> colors = object1.getColors();
-            String string = object1.toString();
+            String string = null != base ? base.format(object1) : object1.toString();
             int cur = 0;
             int count = 0;
             StringBuilder sb = new StringBuilder();
@@ -66,7 +81,7 @@ public static String toColors(final Object object) {
             }
             return sb.toString();
         } else {
-            return object.toString();
+            return null != base ? base.format(object) : object.toString();
         }
     }
 

diff --git a/toolbelt/src/main/java/com/simplifyops/toolbelt/NiceFormatter.java b/toolbelt/src/main/java/com/simplifyops/toolbelt/NiceFormatter.java
@@ -18,6 +18,12 @@ public NiceFormatter(final OutputFormatter base) {
         this.base = base;
     }
 
+    @Override
+    public OutputFormatter withBase(final OutputFormatter base) {
+        this.base = base;
+        return this;
+    }
+
     @Override
     public String format(final Object o) {
         if (o instanceof Map) {

diff --git a/toolbelt/src/main/java/com/simplifyops/toolbelt/OutputFormatter.java b/toolbelt/src/main/java/com/simplifyops/toolbelt/OutputFormatter.java
@@ -5,4 +5,6 @@
  */
 public interface OutputFormatter {
     String format(Object o);
+
+    OutputFormatter withBase(OutputFormatter base);
 }
diff --git a/toolbelt/src/main/java/com/simplifyops/toolbelt/ToStringFormatter.java b/toolbelt/src/main/java/com/simplifyops/toolbelt/ToStringFormatter.java
@@ -8,4 +8,9 @@ public class ToStringFormatter implements OutputFormatter {
     public String format(final Object o) {
         return o != null ? o.toString() : null;
     }
+
+    @Override
+    public OutputFormatter withBase(final OutputFormatter base) {
+        return base;
+    }
 }
diff --git a/toolbelt/src/main/java/com/simplifyops/toolbelt/ToolBelt.java b/toolbelt/src/main/java/com/simplifyops/toolbelt/ToolBelt.java
@@ -15,7 +15,9 @@ public class ToolBelt {
     private CommandInput inputParser;
     private Set<String> helpCommands;
     private CommandOutput commandOutput;
+    private OutputFormatter baseFormatter;
     private OutputFormatter formatter;
+    private boolean ansiColor;
 
     /**
      * Create a simple CLI tool for the object, using {@link SimpleCommandInput} to parse
@@ -94,7 +96,6 @@ public static ToolBelt belt(String name) {
     protected ToolBelt(String name) {
         commands = new CommandSet(name);
         helpCommands = new HashSet<>();
-        formatter = new NiceFormatter(new ToStringFormatter());
     }
 
     /**
@@ -147,8 +148,7 @@ public ToolBelt systemOutput() {
      * @return this builder
      */
     public ToolBelt ansiColorOutput(boolean enabled) {
-        commandOutput(enabled ? new ANSIColorOutput(new SystemOutput()) : new SystemOutput());
-        formatter = new NiceFormatter(enabled ? new ANSIColorOutput(null) : new ToStringFormatter());
+        ansiColor = true;
         return this;
     }
 
@@ -171,6 +171,16 @@ public ToolBelt commandOutput(CommandOutput output) {
         return this;
     }
 
+    /**
+     * Format output data with this formatter
+     *
+     * @return this
+     */
+    public ToolBelt formatter(OutputFormatter formatter) {
+        this.formatter = formatter;
+        return this;
+    }
+
     private static class CommandContext {
         private CommandInput inputParser;
         private CommandOutput output;
@@ -586,8 +596,15 @@ public ToolBelt finalOutput(CommandOutput output) {
         return this;
     }
     public CommandOutput finalOutput() {
+        if (null == commandOutput) {
+            commandOutput(ansiColor ? new ANSIColorOutput(new SystemOutput()) : new SystemOutput());
+        }
+        baseFormatter = new NiceFormatter(ansiColor ? new ANSIColorOutput(null) : new ToStringFormatter());
         if (null == builtOutput) {
-            builtOutput = new FormattedOutput(commandOutput, formatter);
+            builtOutput = new FormattedOutput(
+                    commandOutput,
+                    null != formatter ? formatter.withBase(baseFormatter) : baseFormatter
+            );
         }
         return builtOutput;
     }