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
34namespace benchmark {
35namespace {
36#ifdef BENCHMARK_OS_WINDOWS
37typedef WORD PlatformColorCode;
38#else
39typedef const char* PlatformColorCode;
40#endif
41
42PlatformColorCode 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
85std::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
113std::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
121void 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
128void 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
158bool 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
189