diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h index 00d82dabfcdd5..504de39642e36 100644 --- a/llvm/include/llvm/Support/raw_ostream.h +++ b/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,21 +271,15 @@ 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 @@ -292,11 +287,12 @@ class raw_ostream { 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 diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp index cb586a46cdeeb..bd0a5d4b301f4 100644 --- a/llvm/lib/Support/raw_ostream.cpp +++ b/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(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(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); } diff --git a/llvm/unittests/Support/raw_ostream_test.cpp b/llvm/unittests/Support/raw_ostream_test.cpp index 9f250349b8cbd..267ef195ad0ac 100644 --- a/llvm/unittests/Support/raw_ostream_test.cpp +++ b/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;