Index: llvm/include/llvm/Support/raw_ostream.h =================================================================== --- llvm/include/llvm/Support/raw_ostream.h +++ llvm/include/llvm/Support/raw_ostream.h @@ -64,6 +64,7 @@ /// for a \see write_impl() call to handle the data which has been put into /// this buffer. char *OutBufStart, *OutBufEnd, *OutBufCur; + bool ColorEnabled = true; enum class BufferKind { Unbuffered = 0, @@ -272,19 +273,14 @@ /// @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; - } + 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 @@ -296,7 +292,7 @@ // Enable or disable colors. Once disable_colors() is called, // changeColor() has no effect until enable_colors() is called. - virtual void enable_colors(bool /*enable*/) {} + virtual void enable_colors(bool enable) { ColorEnabled = enable; } //===--------------------------------------------------------------------===// // Subclass Interface @@ -352,6 +348,9 @@ /// unused bytes in the buffer. void copy_to_buffer(const char *Ptr, size_t Size); + /// Compute whether colors should be used. + bool use_color() const; + virtual void anchor(); }; @@ -398,7 +397,6 @@ 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 +462,10 @@ /// 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 @@ -523,8 +513,11 @@ /// currently in the buffer. uint64_t current_pos() const override { return OS.size(); } + bool has_colors() const override { return true; } + public: explicit raw_string_ostream(std::string &O) : OS(O) { + enable_colors(false); SetUnbuffered(); } ~raw_string_ostream() override; Index: llvm/lib/Support/raw_ostream.cpp =================================================================== --- llvm/lib/Support/raw_ostream.cpp +++ llvm/lib/Support/raw_ostream.cpp @@ -499,6 +499,62 @@ return write_padding<'\0'>(*this, NumZeros); } +bool raw_ostream::use_color() const { + // Colors were explicitly disabled. + if (!ColorEnabled) + return false; + + // Colors are not supported by this stream. + if (!has_colors()) + return false; + + // Colors require changing the terminal but this stream is not going to a + // terminal. + if (sys::Process::ColorNeedsFlush() && !is_displayed()) + return false; + + return true; +} + +raw_ostream &raw_ostream::changeColor(enum Colors colors, bool bold, + bool bg) { + if (!use_color()) + return *this; + + if (sys::Process::ColorNeedsFlush()) + flush(); + const char *colorcode = + (colors == SAVEDCOLOR) + ? sys::Process::OutputBold(bg) + : sys::Process::OutputColor(static_cast(colors), bold, bg); + if (colorcode) + write(colorcode, strlen(colorcode)); + return *this; +} + +raw_ostream &raw_ostream::resetColor() { + if (!use_color()) + return *this; + + if (sys::Process::ColorNeedsFlush()) + flush(); + if (const char *colorcode = sys::Process::ResetColor()) + write(colorcode, strlen(colorcode)); + return *this; +} + +raw_ostream &raw_ostream::reverseColor() { + if (!use_color()) + return *this; + + if (sys::Process::ColorNeedsFlush()) + flush(); + if (const char *colorcode = sys::Process::OutputReverse()) + write(colorcode, strlen(colorcode)); + return *this; +} + + void raw_ostream::anchor() {} //===----------------------------------------------------------------------===// @@ -798,58 +854,6 @@ #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(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); } Index: llvm/unittests/Support/raw_ostream_test.cpp =================================================================== --- llvm/unittests/Support/raw_ostream_test.cpp +++ llvm/unittests/Support/raw_ostream_test.cpp @@ -350,6 +350,27 @@ format_bytes_with_ascii_str(B.take_front(12), 0, 7, 1)); } +TEST(raw_ostreamTest, Colors) { + { + std::string S; + raw_string_ostream Sos(S); + Sos.enable_colors(false); + Sos.changeColor(raw_ostream::YELLOW); + EXPECT_TRUE(Sos.str().empty()); + } + + { + std::string S; + raw_string_ostream Sos(S); + Sos.enable_colors(true); + Sos.changeColor(raw_ostream::YELLOW); + EXPECT_FALSE(Sos.str().empty()); +#ifdef LLVM_ON_UNIX + EXPECT_EQ("\x1B[0;33m", Sos.str()); +#endif + } +} + TEST(raw_fd_ostreamTest, multiple_raw_fd_ostream_to_stdout) { std::error_code EC;