[Support] Replace 'DisableColors' boolean with 'ColorMode' enum

Replace the DisableColors with a ColorMode which can be set to Auto,
Enabled and Disabled. The purpose of this change is to make it possible
to ignore the command line option not only for disabling colors, but
also for enabling them.

Differential revision: https://reviews.llvm.org/D81056

include/llvm/Support/WithColor.h

@@ -33,31 +33,43 @@ enum class HighlightColor {
   Remark
 };
 
+enum class ColorMode {
+  /// Determine whether to use color based on the command line argument and the
+  /// raw_ostream.
+  Auto,
+  /// Enable colors. Because raw_ostream is the one implementing colors, this
+  /// has no effect if the stream does not support colors or has colors
+  /// disabled.
+  Enable,
+  /// Disable colors.
+  Disable,
+};
+
 /// An RAII object that temporarily switches an output stream to a specific
 /// color.
 class WithColor {
   raw_ostream &OS;
-  bool DisableColors;
+  ColorMode Mode;
 
 public:
   /// To be used like this: WithColor(OS, HighlightColor::String) << "text";
   /// @param OS The output stream
   /// @param S Symbolic name for syntax element to color
-  /// @param DisableColors Whether to ignore color changes regardless of -color
+  /// @param Mode Enable, disable or compute whether to use colors.
-  /// and support in OS
+  WithColor(raw_ostream &OS, HighlightColor S,
-  WithColor(raw_ostream &OS, HighlightColor S, bool DisableColors = false);
+            ColorMode Mode = ColorMode::Auto);
   /// To be used like this: WithColor(OS, raw_ostream::Black) << "text";
   /// @param OS The output stream
   /// @param Color ANSI color to use, the special SAVEDCOLOR can be used to
   /// change only the bold attribute, and keep colors untouched
   /// @param Bold Bold/brighter text, default false
   /// @param BG If true, change the background, default: change foreground
-  /// @param DisableColors Whether to ignore color changes regardless of -color
+  /// @param Mode Enable, disable or compute whether to use colors.
-  /// and support in OS
   WithColor(raw_ostream &OS,
             raw_ostream::Colors Color = raw_ostream::SAVEDCOLOR,
-            bool Bold = false, bool BG = false, bool DisableColors = false)
+            bool Bold = false, bool BG = false,
-      : OS(OS), DisableColors(DisableColors) {
+            ColorMode Mode = ColorMode::Auto)
+      : OS(OS), Mode(Mode) {
     changeColor(Color, Bold, BG);
   }
   ~WithColor();

lib/Support/SourceMgr.cpp

@@ -450,8 +450,10 @@ static bool isNonASCII(char c) { return c & 0x80; }
 
 void SMDiagnostic::print(const char *ProgName, raw_ostream &OS, bool ShowColors,
                          bool ShowKindLabel) const {
+  ColorMode Mode = ShowColors ? ColorMode::Auto : ColorMode::Disable;
+
   {
-    WithColor S(OS, raw_ostream::SAVEDCOLOR, true, false, !ShowColors);
+    WithColor S(OS, raw_ostream::SAVEDCOLOR, true, false, Mode);
 
     if (ProgName && ProgName[0])
       S << ProgName << ": ";
@@ -488,8 +490,7 @@ void SMDiagnostic::print(const char *ProgName, raw_ostream &OS, bool ShowColors,
     }
   }
 
-  WithColor(OS, raw_ostream::SAVEDCOLOR, true, false, !ShowColors)
+  WithColor(OS, raw_ostream::SAVEDCOLOR, true, false, Mode) << Message << '\n';
-      << Message << '\n';
 
   if (LineNo == -1 || ColumnNo == -1)
     return;
@@ -536,7 +537,8 @@ void SMDiagnostic::print(const char *ProgName, raw_ostream &OS, bool ShowColors,
   printSourceLine(OS, LineContents);
 
   {
-    WithColor S(OS, raw_ostream::GREEN, true, false, !ShowColors);
+    ColorMode Mode = ShowColors ? ColorMode::Auto : ColorMode::Disable;
+    WithColor S(OS, raw_ostream::GREEN, true, false, Mode);
 
     // Print out the caret line, matching tabs in the source line.
     for (unsigned i = 0, e = CaretLine.size(), OutCol = 0; i != e; ++i) {

lib/Support/WithColor.cpp

@@ -18,8 +18,8 @@ static cl::opt<cl::boolOrDefault>
              cl::desc("Use colors in output (default=autodetect)"),
              cl::init(cl::BOU_UNSET));
 
-WithColor::WithColor(raw_ostream &OS, HighlightColor Color, bool DisableColors)
+WithColor::WithColor(raw_ostream &OS, HighlightColor Color, ColorMode Mode)
-    : OS(OS), DisableColors(DisableColors) {
+    : OS(OS), Mode(Mode) {
   // Detect color from terminal type unless the user passed the --color option.
   if (colorsEnabled()) {
     switch (Color) {
@@ -69,7 +69,9 @@ raw_ostream &WithColor::error(raw_ostream &OS, StringRef Prefix,
                               bool DisableColors) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Error, DisableColors).get()
+  return WithColor(OS, HighlightColor::Error,
+                   DisableColors ? ColorMode::Disable : ColorMode::Auto)
+             .get()
          << "error: ";
 }
 
@@ -77,7 +79,9 @@ raw_ostream &WithColor::warning(raw_ostream &OS, StringRef Prefix,
                                 bool DisableColors) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Warning, DisableColors).get()
+  return WithColor(OS, HighlightColor::Warning,
+                   DisableColors ? ColorMode::Disable : ColorMode::Auto)
+             .get()
          << "warning: ";
 }
 
@@ -85,23 +89,33 @@ raw_ostream &WithColor::note(raw_ostream &OS, StringRef Prefix,
                              bool DisableColors) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Note, DisableColors).get() << "note: ";
+  return WithColor(OS, HighlightColor::Note,
+                   DisableColors ? ColorMode::Disable : ColorMode::Auto)
+             .get()
+         << "note: ";
 }
 
 raw_ostream &WithColor::remark(raw_ostream &OS, StringRef Prefix,
                                bool DisableColors) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Remark, DisableColors).get()
+  return WithColor(OS, HighlightColor::Remark,
+                   DisableColors ? ColorMode::Disable : ColorMode::Auto)
+             .get()
          << "remark: ";
 }
 
 bool WithColor::colorsEnabled() {
-  if (DisableColors)
+  switch (Mode) {
+  case ColorMode::Enable:
+    return true;
+  case ColorMode::Disable:
     return false;
-  if (UseColor == cl::BOU_UNSET)
+  case ColorMode::Auto:
-    return OS.has_colors();
+    return UseColor == cl::BOU_UNSET ? OS.has_colors()
-  return UseColor == cl::BOU_TRUE;
+                                     : UseColor == cl::BOU_TRUE;
+  }
+  llvm_unreachable("All cases handled above.");
 }
 
 WithColor &WithColor::changeColor(raw_ostream::Colors Color, bool Bold,

unittests/Support/CMakeLists.txt

@@ -83,6 +83,7 @@ add_llvm_unittest(SupportTests
   UnicodeTest.cpp
   VersionTupleTest.cpp
   VirtualFileSystemTest.cpp
+  WithColorTest.cpp
   YAMLIOTest.cpp
   YAMLParserTest.cpp
   formatted_raw_ostream_test.cpp

unittests/Support/WithColorTest.cpp

@@ -0,0 +1,43 @@
+//===- WithColorTest.cpp --------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/WithColor.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+TEST(WithColorTest, ColorMode) {
+  {
+    std::string S;
+    llvm::raw_string_ostream OS(S);
+    OS.enable_colors(true);
+
+    WithColor(OS, HighlightColor::Error, ColorMode::Disable) << "test";
+    EXPECT_EQ("test", OS.str());
+  }
+
+  {
+    std::string S;
+    llvm::raw_string_ostream OS(S);
+    OS.enable_colors(true);
+
+    WithColor(OS, HighlightColor::Error, ColorMode::Auto) << "test";
+    EXPECT_EQ("test", OS.str());
+  }
+
+#ifdef LLVM_ON_UNIX
+  {
+    std::string S;
+    llvm::raw_string_ostream OS(S);
+    OS.enable_colors(true);
+
+    WithColor(OS, HighlightColor::Error, ColorMode::Enable) << "test";
+    EXPECT_EQ("\x1B[0;1;31mtest\x1B[0m", OS.str());
+  }
+#endif
+}