| 1 | // Copyright 2015 Google Inc. All rights reserved. |
| 2 | // |
| 3 | // Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | // you may not use this file except in compliance with the License. |
| 5 | // You may obtain a copy of the License at |
| 6 | // |
| 7 | // http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | // |
| 9 | // Unless required by applicable law or agreed to in writing, software |
| 10 | // distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | // See the License for the specific language governing permissions and |
| 13 | // limitations under the License. |
| 14 | |
| 15 | #include "colorprint.h" |
| 16 | |
| 17 | #include <cstdarg> |
| 18 | #include <cstdio> |
| 19 | #include <cstdlib> |
| 20 | #include <cstring> |
| 21 | #include <memory> |
| 22 | #include <string> |
| 23 | |
| 24 | #include "check.h" |
| 25 | #include "internal_macros.h" |
| 26 | |
| 27 | #ifdef BENCHMARK_OS_WINDOWS |
| 28 | #include <io.h> |
| 29 | #include <windows.h> |
| 30 | #else |
| 31 | #include <unistd.h> |
| 32 | #endif // BENCHMARK_OS_WINDOWS |
| 33 | |
| 34 | namespace benchmark { |
| 35 | namespace { |
| 36 | #ifdef BENCHMARK_OS_WINDOWS |
| 37 | typedef WORD PlatformColorCode; |
| 38 | #else |
| 39 | typedef const char* PlatformColorCode; |
| 40 | #endif |
| 41 | |
| 42 | PlatformColorCode GetPlatformColorCode(LogColor color) { |
| 43 | #ifdef BENCHMARK_OS_WINDOWS |
| 44 | switch (color) { |
| 45 | case COLOR_RED: |
| 46 | return FOREGROUND_RED; |
| 47 | case COLOR_GREEN: |
| 48 | return FOREGROUND_GREEN; |
| 49 | case COLOR_YELLOW: |
| 50 | return FOREGROUND_RED | FOREGROUND_GREEN; |
| 51 | case COLOR_BLUE: |
| 52 | return FOREGROUND_BLUE; |
| 53 | case COLOR_MAGENTA: |
| 54 | return FOREGROUND_BLUE | FOREGROUND_RED; |
| 55 | case COLOR_CYAN: |
| 56 | return FOREGROUND_BLUE | FOREGROUND_GREEN; |
| 57 | case COLOR_WHITE: // fall through to default |
| 58 | default: |
| 59 | return 0; |
| 60 | } |
| 61 | #else |
| 62 | switch (color) { |
| 63 | case COLOR_RED: |
| 64 | return "1" ; |
| 65 | case COLOR_GREEN: |
| 66 | return "2" ; |
| 67 | case COLOR_YELLOW: |
| 68 | return "3" ; |
| 69 | case COLOR_BLUE: |
| 70 | return "4" ; |
| 71 | case COLOR_MAGENTA: |
| 72 | return "5" ; |
| 73 | case COLOR_CYAN: |
| 74 | return "6" ; |
| 75 | case COLOR_WHITE: |
| 76 | return "7" ; |
| 77 | default: |
| 78 | return nullptr; |
| 79 | }; |
| 80 | #endif |
| 81 | } |
| 82 | |
| 83 | } // end namespace |
| 84 | |
| 85 | std::string FormatString(const char* msg, va_list args) { |
| 86 | // we might need a second shot at this, so pre-emptivly make a copy |
| 87 | va_list args_cp; |
| 88 | va_copy(args_cp, args); |
| 89 | |
| 90 | std::size_t size = 256; |
| 91 | char local_buff[256]; |
| 92 | auto ret = vsnprintf(s: local_buff, maxlen: size, format: msg, arg: args_cp); |
| 93 | |
| 94 | va_end(args_cp); |
| 95 | |
| 96 | // currently there is no error handling for failure, so this is hack. |
| 97 | BM_CHECK(ret >= 0); |
| 98 | |
| 99 | if (ret == 0) { // handle empty expansion |
| 100 | return {}; |
| 101 | } |
| 102 | if (static_cast<size_t>(ret) < size) { |
| 103 | return local_buff; |
| 104 | } |
| 105 | // we did not provide a long enough buffer on our first attempt. |
| 106 | size = static_cast<size_t>(ret) + 1; // + 1 for the null byte |
| 107 | std::unique_ptr<char[]> buff(new char[size]); |
| 108 | ret = vsnprintf(s: buff.get(), maxlen: size, format: msg, arg: args); |
| 109 | BM_CHECK(ret > 0 && (static_cast<size_t>(ret)) < size); |
| 110 | return buff.get(); |
| 111 | } |
| 112 | |
| 113 | std::string FormatString(const char* msg, ...) { |
| 114 | va_list args; |
| 115 | va_start(args, msg); |
| 116 | auto tmp = FormatString(msg, args); |
| 117 | va_end(args); |
| 118 | return tmp; |
| 119 | } |
| 120 | |
| 121 | void ColorPrintf(std::ostream& out, LogColor color, const char* fmt, ...) { |
| 122 | va_list args; |
| 123 | va_start(args, fmt); |
| 124 | ColorPrintf(out, color, fmt, args); |
| 125 | va_end(args); |
| 126 | } |
| 127 | |
| 128 | void ColorPrintf(std::ostream& out, LogColor color, const char* fmt, |
| 129 | va_list args) { |
| 130 | #ifdef BENCHMARK_OS_WINDOWS |
| 131 | ((void)out); // suppress unused warning |
| 132 | |
| 133 | const HANDLE stdout_handle = GetStdHandle(STD_OUTPUT_HANDLE); |
| 134 | |
| 135 | // Gets the current text color. |
| 136 | CONSOLE_SCREEN_BUFFER_INFO buffer_info; |
| 137 | GetConsoleScreenBufferInfo(stdout_handle, &buffer_info); |
| 138 | const WORD old_color_attrs = buffer_info.wAttributes; |
| 139 | |
| 140 | // We need to flush the stream buffers into the console before each |
| 141 | // SetConsoleTextAttribute call lest it affect the text that is already |
| 142 | // printed but has not yet reached the console. |
| 143 | out.flush(); |
| 144 | SetConsoleTextAttribute(stdout_handle, |
| 145 | GetPlatformColorCode(color) | FOREGROUND_INTENSITY); |
| 146 | out << FormatString(fmt, args); |
| 147 | |
| 148 | out.flush(); |
| 149 | // Restores the text color. |
| 150 | SetConsoleTextAttribute(stdout_handle, old_color_attrs); |
| 151 | #else |
| 152 | const char* color_code = GetPlatformColorCode(color); |
| 153 | if (color_code) out << FormatString(msg: "\033[0;3%sm" , color_code); |
| 154 | out << FormatString(msg: fmt, args) << "\033[m" ; |
| 155 | #endif |
| 156 | } |
| 157 | |
| 158 | bool IsColorTerminal() { |
| 159 | #if BENCHMARK_OS_WINDOWS |
| 160 | // On Windows the TERM variable is usually not set, but the |
| 161 | // console there does support colors. |
| 162 | return 0 != _isatty(_fileno(stdout)); |
| 163 | #else |
| 164 | // On non-Windows platforms, we rely on the TERM variable. This list of |
| 165 | // supported TERM values is copied from Google Test: |
| 166 | // <https://github.com/google/googletest/blob/v1.13.0/googletest/src/gtest.cc#L3225-L3259>. |
| 167 | const char* const SUPPORTED_TERM_VALUES[] = { |
| 168 | "xterm" , |
| 169 | "xterm-color" , |
| 170 | "xterm-256color" , |
| 171 | "screen" , |
| 172 | "screen-256color" , |
| 173 | "tmux" , |
| 174 | "tmux-256color" , |
| 175 | "rxvt-unicode" , |
| 176 | "rxvt-unicode-256color" , |
| 177 | "linux" , |
| 178 | "cygwin" , |
| 179 | "xterm-kitty" , |
| 180 | "alacritty" , |
| 181 | "foot" , |
| 182 | "foot-extra" , |
| 183 | "wezterm" , |
| 184 | }; |
| 185 | |
| 186 | const char* const term = getenv(name: "TERM" ); |
| 187 | |
| 188 | bool term_supports_color = false; |
| 189 | for (const char* candidate : SUPPORTED_TERM_VALUES) { |
| 190 | if (term && 0 == strcmp(s1: term, s2: candidate)) { |
| 191 | term_supports_color = true; |
| 192 | break; |
| 193 | } |
| 194 | } |
| 195 | |
| 196 | return 0 != isatty(fd: fileno(stdout)) && term_supports_color; |
| 197 | #endif // BENCHMARK_OS_WINDOWS |
| 198 | } |
| 199 | |
| 200 | } // end namespace benchmark |
| 201 | |