From 5a9dc8f99110ee5ae8d17b9d8e896a1a9560ff4b Mon Sep 17 00:00:00 2001
From: Ryuuke <ryuuke@live.fr>
Date: Sun, 8 Feb 2015 16:08:29 +0000
Subject: [PATCH] windows cli colors revision and test case

---
 format.cc           | 42 ++++++++++++++++++++++++++++++++++++++----
 format.h            | 25 +------------------------
 test/format-test.cc |  9 ++++++++-
 3 files changed, 47 insertions(+), 29 deletions(-)

diff --git a/format.cc b/format.cc
index ba978035..07e630d5 100644
--- a/format.cc
+++ b/format.cc
@@ -35,6 +35,15 @@
 #include <cmath>
 #include <cstdarg>
 
+#ifdef _WIN32
+# define WIN32_LEAN_AND_MEAN
+# ifdef __MINGW32__
+#  include <cstring>
+# endif
+# include <windows.h>
+# undef ERROR
+#endif
+
 using fmt::internal::Arg;
 
 // Check if exceptions are disabled.
@@ -113,9 +122,21 @@ struct IntChecker<true> {
 };
 
 #ifdef _WIN32
-    const char RESET_COLOR = FOREGROUND_GREEN | FOREGROUND_RED | FOREGROUND_BLUE;
+  uint8_t win32_colors[] =
+  {
+    0,
+    FOREGROUND_RED | FOREGROUND_INTENSITY,
+    FOREGROUND_GREEN | FOREGROUND_INTENSITY,
+    FOREGROUND_GREEN | FOREGROUND_RED | FOREGROUND_INTENSITY,
+    FOREGROUND_BLUE | FOREGROUND_INTENSITY,
+    FOREGROUND_BLUE | FOREGROUND_RED | FOREGROUND_INTENSITY,
+    FOREGROUND_INTENSITY | FOREGROUND_GREEN | FOREGROUND_BLUE,
+    FOREGROUND_GREEN | FOREGROUND_RED | FOREGROUND_BLUE | FOREGROUND_INTENSITY
+  };
+  WORD RESET_COLOR = 0;
+  bool reset_color_flag = false;
 #else
-    const char RESET_COLOR[] = "\x1b[0m";
+  const char RESET_COLOR[] = "\x1b[0m";
 #endif
 
 typedef void (*FormatFunc)(fmt::Writer &, int, fmt::StringRef);
@@ -1096,7 +1117,19 @@ FMT_FUNC void fmt::print(std::ostream &os, StringRef format_str, ArgList args) {
 
 FMT_FUNC void fmt::print_colored(Color c, StringRef format, ArgList args) {
 #ifdef _WIN32
-  SetConsoleTextAttribute(GetStdHandle(STD_OUTPUT_HANDLE), c);
+	HANDLE handle = GetStdHandle(STD_OUTPUT_HANDLE);
+	if (handle == INVALID_HANDLE_VALUE)
+		FMT_THROW(GetLastError(), "cannot get output handle");
+	if (!reset_color_flag) {
+		CONSOLE_SCREEN_BUFFER_INFO infoCon;
+		if (!GetConsoleScreenBufferInfo(handle, &infoCon))
+			FMT_THROW(GetLastError(), "cannot get console informations");
+		RESET_COLOR = infoCon.wAttributes;
+		reset_color_flag = true;
+	}
+	WORD color = static_cast<int>(c) >= ARRAYSIZE(win32_colors) ? RESET_COLOR : win32_colors[c];
+	if (!SetConsoleTextAttribute(handle, color))
+		FMT_THROW(GetLastError(), "cannot set console color");
 #else
   char escape[] = "\x1b[30m";
   escape[3] = '0' + static_cast<char>(c);
@@ -1104,7 +1137,8 @@ FMT_FUNC void fmt::print_colored(Color c, StringRef format, ArgList args) {
 #endif
   print(format, args);
 #ifdef _WIN32
-  SetConsoleTextAttribute(GetStdHandle(STD_OUTPUT_HANDLE), RESET_COLOR);
+  if(!SetConsoleTextAttribute(handle, RESET_COLOR))
+    FMT_THROW(GetLastError(), "cannot set console color");
 #else
   std::fputs(RESET_COLOR, stdout);
 #endif
diff --git a/format.h b/format.h
index 1e78f8dd..0af01439 100644
--- a/format.h
+++ b/format.h
@@ -40,15 +40,6 @@
 #include <string>
 #include <sstream>
 
-#ifdef _WIN32
-# define WIN32_LEAN_AND_MEAN
-# ifdef __MINGW32__
-#  include <cstring>
-# endif
-# include <windows.h>
-# undef ERROR
-#endif
-
 #if _SECURE_SCL
 # include <iterator>
 #endif
@@ -2163,21 +2154,7 @@ void report_windows_error(int error_code, StringRef message) FMT_NOEXCEPT(true);
 
 #endif
 
-#ifdef _WIN32
-    enum Color
-    {
-        BLACK = 0,
-        RED = FOREGROUND_RED | FOREGROUND_INTENSITY,
-        GREEN = FOREGROUND_GREEN | FOREGROUND_INTENSITY,
-        YELLOW = FOREGROUND_GREEN | FOREGROUND_RED | FOREGROUND_INTENSITY,
-        BLUE = FOREGROUND_BLUE | FOREGROUND_INTENSITY,
-        MAGENTA = FOREGROUND_BLUE | FOREGROUND_RED | FOREGROUND_INTENSITY,
-        CYAN = FOREGROUND_INTENSITY | FOREGROUND_GREEN | FOREGROUND_BLUE,
-        WHITE = FOREGROUND_GREEN | FOREGROUND_RED | FOREGROUND_BLUE | FOREGROUND_INTENSITY
-    };
-#else
-    enum Color { BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE };
-#endif
+enum Color { BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE };
 
 /**
   Formats a string and prints it to stdout using ANSI escape sequences
diff --git a/test/format-test.cc b/test/format-test.cc
index d1f4f074..e8f968b9 100644
--- a/test/format-test.cc
+++ b/test/format-test.cc
@@ -1371,8 +1371,15 @@ TEST(FormatTest, Print) {
 
 #if FMT_USE_FILE_DESCRIPTORS
 TEST(FormatTest, PrintColored) {
+#if _WIN32
   EXPECT_WRITE(stdout, fmt::print_colored(fmt::RED, "Hello, {}!\n", "world"),
-      "\x1b[31mHello, world!\n\x1b[0m");
+    "Hello, world!\n");
+  EXPECT_WRITE(stdout, fmt::print_colored(static_cast<fmt::Color>(29673),
+    "Hello, {}!\n", "world"), "Hello, world!\n");
+#else
+  EXPECT_WRITE(stdout, fmt::print_colored(fmt::RED, "Hello, {}!\n", "world"),
+    "\x1b[31mHello, world!\n\x1b[0m");
+#endif
 }
 #endif