Refactor CLI printing and colors

config/checkstyle/checkstyle.xml

@@ -163,8 +163,12 @@
         <module name="RightCurly">
             <property name="id" value="RightCurlyAlone"/>
             <property name="option" value="alone"/>
-            <property name="tokens"
-                      value="CLASS_DEF, METHOD_DEF, LITERAL_FOR, LITERAL_WHILE, STATIC_INIT, INSTANCE_INIT"/>
+            <property name="tokens" value="LITERAL_FOR, LITERAL_WHILE, STATIC_INIT, INSTANCE_INIT"/>
+        </module>
+        <module name="RightCurly">
+            <property name="id" value="RightCurlyAloneOrSingle"/>
+            <property name="option" value="alone_or_singleline"/>
+            <property name="tokens" value="CLASS_DEF, METHOD_DEF"/>
         </module>
 
         <!-- Checks for common coding problems               -->

smithy-cli/src/it/java/software/amazon/smithy/cli/IntegUtils.java

@@ -34,6 +34,7 @@
 import java.util.Map;
 import java.util.function.Consumer;
 import software.amazon.smithy.utils.IoUtils;
+import software.amazon.smithy.utils.MapUtils;
 
 public final class IntegUtils {
 
@@ -47,7 +48,7 @@ public static void withProject(String projectName, Consumer<Path> consumer) {
     }
 
     public static void run(String projectName, List<String> args, Consumer<RunResult> consumer) {
-        run(projectName, args, Collections.emptyMap(), consumer);
+        run(projectName, args, MapUtils.of(EnvironmentVariable.NO_COLOR.toString(), "true"), consumer);
     }
 
     public static void run(String projectName, List<String> args, Map<String, String> env,

smithy-cli/src/it/java/software/amazon/smithy/cli/RootCommandTest.java

@@ -18,9 +18,11 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.not;
 
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.utils.ListUtils;
+import software.amazon.smithy.utils.MapUtils;
 
 public class RootCommandTest {
     @Test
@@ -36,6 +38,9 @@ public void passing_h_printsHelp() {
         IntegUtils.run("simple-config-sources", ListUtils.of("-h"), result -> {
             assertThat(result.getExitCode(), equalTo(0));
             ensureHelpOutput(result);
+
+            // We force NO_COLOR by default in the run method. Test that's honored here.
+            assertThat(result.getOutput(), not(containsString("[0m")));
         });
     }
 
@@ -71,6 +76,17 @@ public void errorsOnInvalidArgument() {
         });
     }
 
+    @Test
+    public void runsWithColors() {
+        IntegUtils.run("simple-config-sources",
+                       ListUtils.of("--help"),
+                       MapUtils.of(EnvironmentVariable.FORCE_COLOR.toString(), "true"),
+                       result -> {
+            assertThat(result.getExitCode(), equalTo(0));
+            assertThat(result.getOutput(), containsString("[0m"));
+        });
+    }
+
     private void ensureHelpOutput(RunResult result) {
         // Make sure it's the help output.
         assertThat(result.getOutput(),

smithy-cli/src/main/java/software/amazon/smithy/cli/AnsiColorFormatter.java

@@ -0,0 +1,136 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.cli;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.Objects;
+
+/**
+ * Styles text using ANSI color codes.
+ */
+public enum AnsiColorFormatter implements ColorFormatter {
+
+    /**
+     * Does not write any color.
+     */
+    NO_COLOR {
+        @Override
+        public String style(String text, Style... styles) {
+            return text;
+        }
+
+        @Override
+        public void style(Appendable appendable, String text, Style... styles) {
+            try {
+                appendable.append(text);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        }
+    },
+
+    /**
+     * Writes with ANSI colors.
+     */
+    FORCE_COLOR {
+        @Override
+        public String style(String text, Style... styles) {
+            StringBuilder builder = new StringBuilder();
+            style(builder, text, styles);
+            return builder.toString();
+        }
+
+        @Override
+        public void style(Appendable appendable, String text, Style... styles) {
+            try {
+                if (styles.length == 0) {
+                    appendable.append(text);
+                } else {
+                    appendable.append("\033[");
+                    boolean isAfterFirst = false;
+                    for (Style style : styles) {
+                        if (isAfterFirst) {
+                            appendable.append(';');
+                        }
+                        appendable.append(style.getAnsiColorCode());
+                        isAfterFirst = true;
+                    }
+                    appendable.append('m');
+                    appendable.append(text);
+                    appendable.append("\033[0m");
+                }
+            } catch (IOException e) {
+                throw new CliError("Error writing output", 2, e);
+            }
+        }
+    },
+
+    /**
+     * Writes using ANSI colors if it detects that the environment supports color.
+     */
+    AUTO {
+        private final AnsiColorFormatter delegate = AnsiColorFormatter.detect();
+
+        @Override
+        public String style(String text, Style... styles) {
+            return delegate.style(text, styles);
+        }
+
+        @Override
+        public void style(Appendable appendable, String text, Style... styles) {
+            delegate.style(appendable, text, styles);
+        }
+    };
+
+    /**
+     * Detects if ANSI colors are supported and returns the appropriate Ansi enum variant.
+     *
+     * <p>This method differs from using the {@link AnsiColorFormatter#AUTO} variant directly because it will
+     * detect any changes to the environment that might enable or disable colors.
+     *
+     * @return Returns the detected ANSI color enum variant.
+     */
+    public static AnsiColorFormatter detect() {
+        return isAnsiEnabled() ? FORCE_COLOR : NO_COLOR;
+    }
+
+    private static boolean isAnsiEnabled() {
+        if (EnvironmentVariable.FORCE_COLOR.isSet()) {
+            return true;
+        }
+
+        // Disable colors if NO_COLOR is set to anything.
+        if (EnvironmentVariable.NO_COLOR.isSet()) {
+            return false;
+        }
+
+        String term = EnvironmentVariable.TERM.get();
+
+        // If term is set to "dumb", then don't use colors.
+        if (Objects.equals(term, "dumb")) {
+            return false;
+        }
+
+        // If TERM isn't set at all and Windows is detected, then don't use colors.
+        if (term == null && System.getProperty("os.name").contains("win")) {
+            return false;
+        }
+
+        // Disable colors if no console is associated.
+        return System.console() != null;
+    }
+}

smithy-cli/src/main/java/software/amazon/smithy/cli/ArgumentReceiver.java

@@ -61,7 +61,5 @@ default Consumer<String> testParameter(String name) {
      *
      * @param printer Printer to modify.
      */
-    default void registerHelp(HelpPrinter printer) {
-        // do nothing by default.
-    }
+    default void registerHelp(HelpPrinter printer) {}
 }

smithy-cli/src/main/java/software/amazon/smithy/cli/Cli.java

@@ -15,7 +15,10 @@
 
 package software.amazon.smithy.cli;
 
+import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.util.Arrays;
+import java.util.function.Supplier;
 import java.util.logging.Logger;
 
 /**
@@ -31,14 +34,12 @@ public final class Cli {
 
     private static final Logger LOGGER = Logger.getLogger(Cli.class.getName());
 
-    // Delegate to the stdout consumer by default since this can change.
-    private CliPrinter stdoutPrinter = new CliPrinter.ConsumerPrinter(str -> System.out.print(str));
-
-    // Don't use a method reference in case System.err is changed after initialization.
-    private CliPrinter stdErrPrinter = new CliPrinter.ConsumerPrinter(str -> System.err.print(str));
-
+    private ColorFormatter colorFormatter;
+    private CliPrinter stdoutPrinter;
+    private CliPrinter stderrPrinter;
     private final ClassLoader classLoader;
     private final Command command;
+    private final StandardOptions standardOptions = new StandardOptions();
 
     /**
      * Creates a new CLI with the given name.
@@ -60,40 +61,87 @@ public Cli(Command command, ClassLoader classLoader) {
      */
     public int run(String[] args) {
         Arguments arguments = new Arguments(args);
-        StandardOptions standardOptions = new StandardOptions();
         arguments.addReceiver(standardOptions);
 
-        // Use or disable ANSI escapes in the printers.
-        CliPrinter out = new CliPrinter.ColorPrinter(stdoutPrinter, standardOptions);
-        CliPrinter err = new CliPrinter.ColorPrinter(stdErrPrinter, standardOptions);
+        if (colorFormatter == null) {
+            // CLI arguments haven't been parsed yet, so the CLI doesn't know if --force-color or --no-color
+            // was passed. Defer the color setting implementation by asking StandardOptions before each write.
+            colorFormatter(createDelegatedColorFormatter(standardOptions::colorSetting));
+        }
+
+        if (stdoutPrinter == null) {
+            stdout(CliPrinter.fromOutputStream(System.out));
+        }
+
+        if (stderrPrinter == null) {
+            stderr(CliPrinter.fromOutputStream(System.err));
+        }
 
         // Setup logging after parsing all arguments.
         arguments.onComplete((opts, positional) -> {
-            LoggingUtil.configureLogging(opts.getReceiver(StandardOptions.class), err);
+            LoggingUtil.configureLogging(opts.getReceiver(StandardOptions.class), colorFormatter, stderrPrinter);
             LOGGER.fine(() -> "Running CLI command: " + Arrays.toString(args));
         });
 
         try {
-            return command.execute(arguments, new Command.Env(out, err, classLoader));
-        } catch (Exception e) {
-            err.printException(e, standardOptions.stackTrace());
-            throw CliError.wrap(e);
-        } finally {
             try {
-                LoggingUtil.restoreLogging();
-            } catch (RuntimeException e) {
-                // Show the error, but don't fail the CLI since most invocations are one-time use.
-                err.println(err.style("Unable to restore logging to previous settings", Style.RED));
-                err.printException(e, standardOptions.stackTrace());
+                Command.Env env = new Command.Env(colorFormatter, stdoutPrinter, stderrPrinter, classLoader);
+                return command.execute(arguments, env);
+            } catch (Exception e) {
+                printException(e, standardOptions.stackTrace());
+                throw CliError.wrap(e);
+            } finally {
+                try {
+                    LoggingUtil.restoreLogging();
+                } catch (RuntimeException e) {
+                    // Show the error, but don't fail the CLI since most invocations are one-time use.
+                    printException(e, standardOptions.stackTrace());
+                }
             }
+        } finally {
+            stdoutPrinter.flush();
+            stderrPrinter.flush();
         }
     }
 
-    public void stdout(CliPrinter printer) {
-        stdoutPrinter = printer;
+    public void colorFormatter(ColorFormatter colorFormatter) {
+        this.colorFormatter = colorFormatter;
+    }
+
+    public void stdout(CliPrinter stdoutPrinter) {
+        this.stdoutPrinter = stdoutPrinter;
+    }
+
+    public void stderr(CliPrinter stderrPrinter) {
+        this.stderrPrinter = stderrPrinter;
+    }
+
+    private void printException(Throwable e, boolean stacktrace) {
+        if (!stacktrace) {
+            colorFormatter.println(stderrPrinter, e.getMessage(), Style.RED);
+        } else {
+            try (ColorFormatter.PrinterBuffer buffer = colorFormatter.printerBuffer(stderrPrinter)) {
+                StringWriter writer = new StringWriter();
+                e.printStackTrace(new PrintWriter(writer));
+                String result = writer.toString();
+                int positionOfName = result.indexOf(':');
+                buffer.print(result.substring(0, positionOfName), Style.RED, Style.UNDERLINE);
+                buffer.println(result.substring(positionOfName));
+            }
+        }
     }
 
-    public void stderr(CliPrinter printer) {
-        stdErrPrinter = printer;
+    private static ColorFormatter createDelegatedColorFormatter(Supplier<ColorFormatter> delegateSupplier) {
+        return new ColorFormatter() {
+            @Override
+            public String style(String text, Style... styles) {
+                return delegateSupplier.get().style(text, styles);
+            }
+
+            @Override
+            public void style(Appendable appendable, String text, Style... styles) {
+                delegateSupplier.get().style(appendable, text, styles);
+            }
+        };
     }
 }