2 files changed, 39 insertions, 9 deletions

diff --git a/llvm/lib/Support/WithColor.cpp b/llvm/lib/Support/WithColor.cpp
index 345dd9cf394..1c380c65902 100644
--- a/llvm/lib/Support/WithColor.cpp
+++ b/llvm/lib/Support/WithColor.cpp
@@ -22,6 +22,8 @@ WithColor::WithColor(raw_ostream &OS, HighlightColor Color, bool DisableColors)
     : OS(OS), DisableColors(DisableColors) {
   // Detect color from terminal type unless the user passed the --color option.
   if (colorsEnabled()) {
+    OS.enable_colors();
+
     switch (Color) {
     case HighlightColor::Address:
       OS.changeColor(raw_ostream::YELLOW);
@@ -104,10 +106,9 @@ bool WithColor::colorsEnabled() {
   return UseColor == cl::BOU_TRUE;
 }
 
-WithColor &WithColor::changeColor(raw_ostream::Colors Color, bool Bold,
-                                  bool BG) {
+WithColor &WithColor::changeColor(raw_ostream::Color C, bool Bold, bool BG) {
   if (colorsEnabled())
-    OS.changeColor(Color, Bold, BG);
+    OS.changeColor(C, Bold, BG);
   return *this;
 }
 
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 2baccaa0cbd..da24fe29471 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -65,6 +65,17 @@
 
 using namespace llvm;
 
+const raw_ostream::Color raw_ostream::BLACK;
+const raw_ostream::Color raw_ostream::RED;
+const raw_ostream::Color raw_ostream::GREEN;
+const raw_ostream::Color raw_ostream::YELLOW;
+const raw_ostream::Color raw_ostream::BLUE;
+const raw_ostream::Color raw_ostream::MAGENTA;
+const raw_ostream::Color raw_ostream::CYAN;
+const raw_ostream::Color raw_ostream::WHITE;
+const raw_ostream::Color raw_ostream::SAVEDCOLOR;
+const raw_ostream::Color raw_ostream::RESET;
+
 raw_ostream::~raw_ostream() {
   // raw_ostream's subclasses should take care to flush the buffer
   // in their destructors.
@@ -133,6 +144,14 @@ raw_ostream &raw_ostream::write_hex(unsigned long long N) {
   return *this;
 }
 
+raw_ostream &raw_ostream::operator<<(Color C) {
+  if (C == Color::RESET)
+    resetColor();
+  else
+    changeColor(C);
+  return *this;
+}
+
 raw_ostream &raw_ostream::write_uuid(const uuid_t UUID) {
   for (int Idx = 0; Idx < 16; ++Idx) {
     *this << format("%02" PRIX32, UUID[Idx]);
@@ -552,8 +571,9 @@ raw_fd_ostream::raw_fd_ostream(StringRef Filename, std::error_code &EC,
 /// FD is the file descriptor that this writes to.  If ShouldClose is true, this
 /// closes the file when the stream is destroyed.
 raw_fd_ostream::raw_fd_ostream(int fd, bool shouldClose, bool unbuffered)
-    : raw_pwrite_stream(unbuffered), FD(fd), ShouldClose(shouldClose) {
-  if (FD < 0 ) {
+    : raw_pwrite_stream(unbuffered), FD(fd), ShouldClose(shouldClose),
+      ColorEnabled(sys::Process::FileDescriptorHasColors(fd)) {
+  if (FD < 0) {
     ShouldClose = false;
     return;
   }
@@ -782,13 +802,16 @@ size_t raw_fd_ostream::preferred_buffer_size() const {
 #endif
 }
 
-raw_ostream &raw_fd_ostream::changeColor(enum Colors colors, bool bold,
-                                         bool bg) {
+raw_ostream &raw_fd_ostream::changeColor(Color color, 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(colors, bold, bg);
+      (color == Color::SAVEDCOLOR)
+          ? sys::Process::OutputBold(bg)
+          : sys::Process::OutputColor(static_cast<char>(color), bold, bg);
   if (colorcode) {
     size_t len = strlen(colorcode);
     write(colorcode, len);
@@ -799,6 +822,9 @@ raw_ostream &raw_fd_ostream::changeColor(enum Colors colors, bool bold,
 }
 
 raw_ostream &raw_fd_ostream::resetColor() {
+  if (!ColorEnabled)
+    return *this;
+
   if (sys::Process::ColorNeedsFlush())
     flush();
   const char *colorcode = sys::Process::ResetColor();
@@ -812,6 +838,9 @@ raw_ostream &raw_fd_ostream::resetColor() {
 }
 
 raw_ostream &raw_fd_ostream::reverseColor() {
+  if (!ColorEnabled)
+    return *this;
+
   if (sys::Process::ColorNeedsFlush())
     flush();
   const char *colorcode = sys::Process::OutputReverse();