[Support] Move color handling from raw_fd_ostream to raw_ostream

Move the color handling code from raw_fd_ostream to raw_ostream. This
makes it possible to use colors with any ostream when enabled. The
existing behavior where only raw_fd_ostream supports colors by default
remains unchanged.

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

llvm/include/llvm/Support/raw_ostream.h

@@ -64,6 +64,7 @@ class raw_ostream {
   /// for a \see write_impl() call to handle the data which has been put into
   /// this buffer.
   char *OutBufStart, *OutBufEnd, *OutBufCur;
+  bool ColorEnabled = false;
 
   enum class BufferKind {
     Unbuffered = 0,
@@ -270,33 +271,28 @@ class raw_ostream {
   /// @param Bold bold/brighter text, default false
   /// @param BG if true change the background, default: change foreground
   /// @returns itself so it can be used within << invocations
-  virtual raw_ostream &changeColor(enum Colors Color,
-                                   bool Bold = false,
-                                   bool BG = false) {
-    (void)Color;
-    (void)Bold;
-    (void)BG;
-    return *this;
-  }
+  virtual raw_ostream &changeColor(enum Colors Color, bool Bold = false,
+                                   bool BG = false);
 
   /// Resets the colors to terminal defaults. Call this when you are done
   /// outputting colored text, or before program exit.
-  virtual raw_ostream &resetColor() { return *this; }
+  virtual raw_ostream &resetColor();
 
   /// Reverses the foreground and background colors.
-  virtual raw_ostream &reverseColor() { return *this; }
+  virtual raw_ostream &reverseColor();
 
   /// This function determines if this stream is connected to a "tty" or
   /// "console" window. That is, the output would be displayed to the user
   /// rather than being put on a pipe or stored in a file.
   virtual bool is_displayed() const { return false; }
 
   /// This function determines if this stream is displayed and supports colors.
+  /// The result is unaffected by calls to enable_color().
   virtual bool has_colors() const { return is_displayed(); }
 
-  // Enable or disable colors. Once disable_colors() is called,
-  // changeColor() has no effect until enable_colors() is called.
-  virtual void enable_colors(bool /*enable*/) {}
+  // Enable or disable colors. Once enable_colors(false) is called,
+  // changeColor() has no effect until enable_colors(true) is called.
+  virtual void enable_colors(bool enable) { ColorEnabled = enable; }
 
   //===--------------------------------------------------------------------===//
   // Subclass Interface
@@ -352,6 +348,10 @@ class raw_ostream {
   /// unused bytes in the buffer.
   void copy_to_buffer(const char *Ptr, size_t Size);
 
+  /// Compute whether colors should be used and do the necessary work such as
+  /// flushing. The result is affected by calls to enable_color().
+  bool prepare_colors();
+
   virtual void anchor();
 };
 
@@ -398,7 +398,6 @@ class raw_fd_ostream : public raw_pwrite_stream {
   int FD;
   bool ShouldClose;
   bool SupportsSeeking = false;
-  bool ColorEnabled = true;
 
 #ifdef _WIN32
   /// True if this fd refers to a Windows console device. Mintty and other
@@ -464,18 +463,10 @@ class raw_fd_ostream : public raw_pwrite_stream {
   /// to the offset specified from the beginning of the file.
   uint64_t seek(uint64_t off);
 
-  raw_ostream &changeColor(enum Colors colors, bool bold=false,
-                           bool bg=false) override;
-  raw_ostream &resetColor() override;
-
-  raw_ostream &reverseColor() override;
-
   bool is_displayed() const override;
 
   bool has_colors() const override;
 
-  void enable_colors(bool enable) override { ColorEnabled = enable; }
-
   std::error_code error() const { return EC; }
 
   /// Return the value of the flag in this raw_fd_ostream indicating whether an

llvm/lib/Support/raw_ostream.cpp

@@ -499,6 +499,53 @@ raw_ostream &raw_ostream::write_zeros(unsigned NumZeros) {
   return write_padding<'\0'>(*this, NumZeros);
 }
 
+bool raw_ostream::prepare_colors() {
+  // Colors were explicitly disabled.
+  if (!ColorEnabled)
+    return false;
+
+  // Colors require changing the terminal but this stream is not going to a
+  // terminal.
+  if (sys::Process::ColorNeedsFlush() && !is_displayed())
+    return false;
+
+  if (sys::Process::ColorNeedsFlush())
+    flush();
+
+  return true;
+}
+
+raw_ostream &raw_ostream::changeColor(enum Colors colors, bool bold, bool bg) {
+  if (!prepare_colors())
+    return *this;
+
+  const char *colorcode =
+      (colors == SAVEDCOLOR)
+          ? sys::Process::OutputBold(bg)
+          : sys::Process::OutputColor(static_cast<char>(colors), bold, bg);
+  if (colorcode)
+    write(colorcode, strlen(colorcode));
+  return *this;
+}
+
+raw_ostream &raw_ostream::resetColor() {
+  if (!prepare_colors())
+    return *this;
+
+  if (const char *colorcode = sys::Process::ResetColor())
+    write(colorcode, strlen(colorcode));
+  return *this;
+}
+
+raw_ostream &raw_ostream::reverseColor() {
+  if (!prepare_colors())
+    return *this;
+
+  if (const char *colorcode = sys::Process::OutputReverse())
+    write(colorcode, strlen(colorcode));
+  return *this;
+}
+
 void raw_ostream::anchor() {}
 
 //===----------------------------------------------------------------------===//
@@ -574,6 +621,8 @@ raw_fd_ostream::raw_fd_ostream(int fd, bool shouldClose, bool unbuffered)
     return;
   }
 
+  enable_colors(true);
+
   // Do not attempt to close stdout or stderr. We used to try to maintain the
   // property that tools that support writing file to stdout should not also
   // write informational output to stdout, but in practice we were never able to
@@ -798,58 +847,6 @@ size_t raw_fd_ostream::preferred_buffer_size() const {
 #endif
 }
 
-raw_ostream &raw_fd_ostream::changeColor(enum Colors colors, bool bold,
-                                         bool bg) {
-  if (!ColorEnabled)
-    return *this;
-
-  if (sys::Process::ColorNeedsFlush())
-    flush();
-  const char *colorcode =
-      (colors == SAVEDCOLOR)
-          ? sys::Process::OutputBold(bg)
-          : sys::Process::OutputColor(static_cast<char>(colors), bold, bg);
-  if (colorcode) {
-    size_t len = strlen(colorcode);
-    write(colorcode, len);
-    // don't account colors towards output characters
-    pos -= len;
-  }
-  return *this;
-}
-
-raw_ostream &raw_fd_ostream::resetColor() {
-  if (!ColorEnabled)
-    return *this;
-
-  if (sys::Process::ColorNeedsFlush())
-    flush();
-  const char *colorcode = sys::Process::ResetColor();
-  if (colorcode) {
-    size_t len = strlen(colorcode);
-    write(colorcode, len);
-    // don't account colors towards output characters
-    pos -= len;
-  }
-  return *this;
-}
-
-raw_ostream &raw_fd_ostream::reverseColor() {
-  if (!ColorEnabled)
-    return *this;
-
-  if (sys::Process::ColorNeedsFlush())
-    flush();
-  const char *colorcode = sys::Process::OutputReverse();
-  if (colorcode) {
-    size_t len = strlen(colorcode);
-    write(colorcode, len);
-    // don't account colors towards output characters
-    pos -= len;
-  }
-  return *this;
-}
-
 bool raw_fd_ostream::is_displayed() const {
   return sys::Process::FileDescriptorIsDisplayed(FD);
 }

llvm/unittests/Support/raw_ostream_test.cpp

@@ -350,6 +350,26 @@ TEST(raw_ostreamTest, FormattedHexBytes) {
             format_bytes_with_ascii_str(B.take_front(12), 0, 7, 1));
 }
 
+#ifdef LLVM_ON_UNIX
+TEST(raw_ostreamTest, Colors) {
+  {
+    std::string S;
+    raw_string_ostream Sos(S);
+    Sos.enable_colors(false);
+    Sos.changeColor(raw_ostream::YELLOW);
+    EXPECT_EQ("", Sos.str());
+  }
+
+  {
+    std::string S;
+    raw_string_ostream Sos(S);
+    Sos.enable_colors(true);
+    Sos.changeColor(raw_ostream::YELLOW);
+    EXPECT_EQ("\x1B[0;33m", Sos.str());
+  }
+}
+#endif
+
 TEST(raw_fd_ostreamTest, multiple_raw_fd_ostream_to_stdout) {
   std::error_code EC;