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

[JDevlieghere]

Summary

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 raw_ostream::has_colors() not only for disabling colors, but also for enabling them.

[jhenderson]

Seems reasonable. Should there be some kind of unit-testing involved?

It might be nice (not now but in the future) to change the interface of the WithColor constructors, to change the DisableColors boolean into the ColorMode enum too.

[MaskRay]

In D81056#2070328, @jhenderson wrote:

Seems reasonable. Should there be some kind of unit-testing involved?

It is a good idea. WithColor does not have unit tests yet (llvm/unittests/Support/WithColorTest.cpp)

It might be nice (not now but in the future) to change the interface of the WithColor constructors, to change the DisableColors boolean into the ColorMode enum too.

[JDevlieghere]

Add unit test

[JDevlieghere]

llvm/unittests/Support/WithColorTest.cpp

40

Not sure if this works on Windows, I'll have to double check.

[MaskRay]

LGTM. Hope harbomaster will tell whether this works on Windows.

[jhenderson]

LGTM, assuming the unit test works.

llvm/unittests/Support/WithColorTest.cpp

40

I tried checking, but had trouble applying the parent patch for some reason. Not sure why, sorry.

[JDevlieghere]

With the new approach this change isn't strictly necessary for D81058 to work, but I still think it's an (small) improvement over the status quo. I'm happy to either land or abandon the patch, whatever the community feels is most desirable.

[jdenny]

llvm/include/llvm/Support/WithColor.h

57–58

Please update the comments in this file.

[jhenderson]

In D81056#2076778, @JDevlieghere wrote:

With the new approach this change isn't strictly necessary for D81058 to work, but I still think it's an (small) improvement over the status quo. I'm happy to either land or abandon the patch, whatever the community feels is most desirable.

I'm always a fan of replacing bool arguments with a useful enum, so I'm happy if you think there's still some use in the patch.

llvm/include/llvm/Support/WithColor.h

enum class HighlightColor {
   Enumerator,
   Macro,
   Error,
   Warning,
   Note,
   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
-  /// and support in OS
-  WithColor(raw_ostream &OS, HighlightColor S, bool DisableColors = false);
+  /// @param Mode Enable, disable or compute whether to use colors.
+  WithColor(raw_ostream &OS, HighlightColor S,
+            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
-  /// and support in OS
+  /// @param Mode Enable, disable or compute whether to use colors.
   WithColor(raw_ostream &OS,
             raw_ostream::Colors Color = raw_ostream::SAVEDCOLOR,
-            bool Bold = false, bool BG = false, bool DisableColors = false)
-      : OS(OS), DisableColors(DisableColors) {
+            bool Bold = false, bool BG = false,
+            ColorMode Mode = ColorMode::Auto)
+      : OS(OS), Mode(Mode) {
     changeColor(Color, Bold, BG);
   }
   ~WithColor();
 
   raw_ostream &get() { return OS; }
   operator raw_ostream &() { return OS; }
   template <typename T> WithColor &operator<<(T &O) {
     OS << O;

llvm/lib/Support/SourceMgr.cpp

static void printSourceLine(raw_ostream &S, StringRef LineContents) {
   }
   S << '\n';
 }
 
 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 << ": ";
 
     if (!Filename.empty()) {
       if (Filename == "-")
         S << "<stdin>";
       else
case SourceMgr::DK_Note:
       WithColor::note(OS, "", !ShowColors);
       break;
     case SourceMgr::DK_Remark:
       WithColor::remark(OS, "", !ShowColors);
       break;
     }
   }
 
-  WithColor(OS, raw_ostream::SAVEDCOLOR, true, false, !ShowColors)
-      << Message << '\n';
+  WithColor(OS, raw_ostream::SAVEDCOLOR, true, false, Mode) << Message << '\n';
 
   if (LineNo == -1 || ColumnNo == -1)
     return;
 
   // FIXME: If there are multibyte or multi-column characters in the source, all
   // our ranges will be wrong. To do this properly, we'll need a byte-to-column
   // map like Clang's TextDiagnostic. For now, we'll just handle tabs by
   // expanding them later, and bail out rather than show incorrect ranges and
void SMDiagnostic::print(const char *ProgName, raw_ostream &OS, bool ShowColors,
   // ... and remove trailing whitespace so the output doesn't wrap for it.  We
   // know that the line isn't completely empty because it has the caret in it at
   // least.
   CaretLine.erase(CaretLine.find_last_not_of(' ') + 1);
 
   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) {
       if (i >= LineContents.size() || LineContents[i] != '\t') {
         S << CaretLine[i];
         ++OutCol;
         continue;
       }

llvm/lib/Support/WithColor.cpp

 
 cl::OptionCategory llvm::ColorCategory("Color Options");
 
 static cl::opt<cl::boolOrDefault>
     UseColor("color", cl::cat(ColorCategory),
              cl::desc("Use colors in output (default=autodetect)"),
              cl::init(cl::BOU_UNSET));
 
-WithColor::WithColor(raw_ostream &OS, HighlightColor Color, bool DisableColors)
-    : OS(OS), DisableColors(DisableColors) {
+WithColor::WithColor(raw_ostream &OS, HighlightColor Color, ColorMode Mode)
+    : OS(OS), Mode(Mode) {
   // Detect color from terminal type unless the user passed the --color option.
   if (colorsEnabled()) {
     switch (Color) {
     case HighlightColor::Address:
       OS.changeColor(raw_ostream::YELLOW);
       break;
     case HighlightColor::String:
       OS.changeColor(raw_ostream::GREEN);

 raw_ostream &WithColor::note() { return note(errs()); }
 
 raw_ostream &WithColor::remark() { return remark(errs()); }
 
 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: ";
 }
 
 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: ";
 }
 
 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)
-    return OS.has_colors();
-  return UseColor == cl::BOU_TRUE;
+  case ColorMode::Auto:
+    return UseColor == cl::BOU_UNSET ? OS.has_colors()
+                                     : UseColor == cl::BOU_TRUE;
+  }
+  llvm_unreachable("All cases handled above.");
 }
 
 WithColor &WithColor::changeColor(raw_ostream::Colors Color, bool Bold,
                                   bool BG) {
   if (colorsEnabled())
     OS.changeColor(Color, Bold, BG);
   return *this;
 }

llvm/unittests/Support/CMakeLists.txt

add_llvm_unittest(SupportTests
   ToolOutputFileTest.cpp
   TypeNameTest.cpp
   TypeTraitsTest.cpp
   TrailingObjectsTest.cpp
   TrigramIndexTest.cpp
   UnicodeTest.cpp
   VersionTupleTest.cpp
   VirtualFileSystemTest.cpp
+  WithColorTest.cpp
   YAMLIOTest.cpp
   YAMLParserTest.cpp
   formatted_raw_ostream_test.cpp
   raw_ostream_test.cpp
   raw_pwrite_stream_test.cpp
   raw_sha1_ostream_test.cpp
   xxhashTest.cpp
   )

llvm/unittests/Support/WithColorTest.cpp

+//===- 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
+}