utils/benchmark/src/colorprint.cc

   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 <windows.h>
  29 #include <io.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(local_buff, size, msg, args_cp);
  93
  94   va_end(args_cp);
  95
  96   // currently there is no error handling for failure, so this is hack.
  97   CHECK(ret >= 0);
  98
  99   if (ret == 0)  // handle empty expansion
 100     return {};
 101   else if (static_cast<size_t>(ret) < size)
 102     return local_buff;
 103   else {
 104     // we did not provide a long enough buffer on our first attempt.
 105     size = (size_t)ret + 1;  // + 1 for the null byte
 106     std::unique_ptr<char[]> buff(new char[size]);
 107     ret = vsnprintf(buff.get(), size, msg, args);
 108     CHECK(ret > 0 && ((size_t)ret) < size);
 109     return buff.get();
 110   }
 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   fflush(stdout);
 144   SetConsoleTextAttribute(stdout_handle,
 145                           GetPlatformColorCode(color) | FOREGROUND_INTENSITY);
 146   vprintf(fmt, args);
 147
 148   fflush(stdout);
 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("\033[0;3%sm", color_code);
 154   out << FormatString(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/master/googletest/src/gtest.cc#L2925>.
 167   const char* const SUPPORTED_TERM_VALUES[] = {
 168       "xterm",         "xterm-color",     "xterm-256color",
 169       "screen",        "screen-256color", "tmux",
 170       "tmux-256color", "rxvt-unicode",    "rxvt-unicode-256color",
 171       "linux",         "cygwin",
 172   };
 173
 174   const char* const term = getenv("TERM");
 175
 176   bool term_supports_color = false;
 177   for (const char* candidate : SUPPORTED_TERM_VALUES) {
 178     if (term && 0 == strcmp(term, candidate)) {
 179       term_supports_color = true;
 180       break;
 181     }
 182   }
 183
 184   return 0 != isatty(fileno(stdout)) && term_supports_color;
 185 #endif  // BENCHMARK_OS_WINDOWS
 186 }
 187
 188 }  // end namespace benchmark