Implement basic color support for the Windows console

Only \033[*m (SGR) is supported, with a 16+16 color terminal.

1 files changed, 192 insertions, 37 deletions

diff --git a/src/WinConsoleFile.cc b/src/WinConsoleFile.cc
index 157a2e36..6ef59d25 100644
--- a/src/WinConsoleFile.cc
+++ b/src/WinConsoleFile.cc
@@ -32,44 +32,98 @@
  * files in the program, then also delete it here.
  */
 /* copyright --> */
+
 #include "WinConsoleFile.h"
 
 #include <cstring>
 #include <cstdio>
 #include <cstdarg>
-#include <string>
+#include <vector>
 
 #include "a2io.h"
 #include "util.h"
 
+namespace {
+
+#define FOREGROUND_BLACK 0
+#define FOREGROUND_WHITE FOREGROUND_RED|FOREGROUND_GREEN|FOREGROUND_BLUE
+
+#define BACKGROUND_BLACK 0
+#define BACKGROUND_WHITE BACKGROUND_RED|BACKGROUND_GREEN|BACKGROUND_BLUE
+
+const WORD kForeground[] = {
+  FOREGROUND_BLACK,                   // black
+  FOREGROUND_RED,                     // red
+  FOREGROUND_GREEN,                   // green
+  FOREGROUND_RED | FOREGROUND_GREEN,  // yellow
+  FOREGROUND_BLUE,                    // blue
+  FOREGROUND_BLUE | FOREGROUND_RED,   // magenta
+  FOREGROUND_BLUE | FOREGROUND_GREEN, // cyan
+  FOREGROUND_WHITE                    // white
+};
+
+const WORD kBackground[] = {
+  BACKGROUND_BLACK,                   // black
+  BACKGROUND_RED,                     // red
+  BACKGROUND_GREEN,                   // green
+  BACKGROUND_RED | BACKGROUND_GREEN,  // yellow
+  BACKGROUND_BLUE,                    // blue
+  BACKGROUND_BLUE | BACKGROUND_RED,   // magenta
+  BACKGROUND_BLUE | BACKGROUND_GREEN, // cyan
+  BACKGROUND_WHITE                    // white
+};
+
+} // namespace
+
 namespace aria2 {
 
 WinConsoleFile::WinConsoleFile(DWORD stdHandle)
-  : stdHandle_(stdHandle)
-{}
+  : stdHandle_(stdHandle), bold_(false), underline_(false), reverse_(false),
+    fg_(FOREGROUND_WHITE), bg_(BACKGROUND_BLACK)
+{
+  if (supportsColor()) {
+    CONSOLE_SCREEN_BUFFER_INFO info;
+    GetConsoleScreenBufferInfo(handle(), &info);
+    info.wAttributes &= ~(COMMON_LVB_LEADING_BYTE |
+                          COMMON_LVB_TRAILING_BYTE |
+                          COMMON_LVB_GRID_HORIZONTAL |
+                          COMMON_LVB_GRID_LVERTICAL |
+                          COMMON_LVB_GRID_RVERTICAL |
+                          COMMON_LVB_REVERSE_VIDEO |
+                          COMMON_LVB_UNDERSCORE);
+    fg_ = info.wAttributes & ~(BACKGROUND_BLUE |
+                               BACKGROUND_GREEN |
+                               BACKGROUND_RED |
+                               BACKGROUND_INTENSITY);
+    bg_ = info.wAttributes & ~(FOREGROUND_BLUE |
+                               FOREGROUND_GREEN |
+                               FOREGROUND_RED |
+                               FOREGROUND_INTENSITY);
+    bg_ = (bg_ >>  4) & 0x0F;
+    bold_ = info.wAttributes & FOREGROUND_INTENSITY;
+    underline_ = info.wAttributes & BACKGROUND_INTENSITY;
+  }
 
-WinConsoleFile::~WinConsoleFile() {}
+  deffg_ = fg_;
+  defbg_ = bg_;
+}
 
-namespace {
-bool console(DWORD stdHandle)
+bool WinConsoleFile::supportsColor()
 {
   DWORD mode;
-  return GetConsoleMode(GetStdHandle(stdHandle), &mode);
+  return GetConsoleMode(handle(), &mode);
 }
-} // namespace
 
 size_t WinConsoleFile::write(const char* str)
 {
-  DWORD written;
-  if(console(stdHandle_)) {
-    std::wstring msg = utf8ToWChar(str);
-    WriteConsoleW(GetStdHandle(stdHandle_),
-                  msg.c_str(), msg.size(), &written, 0);
-  } else {
-    WriteFile(GetStdHandle(stdHandle_),
-              str, strlen(str), &written, 0);
+  if (!supportsColor()) {
+    DWORD written = 0;
+    WriteFile(handle(), str, strlen(str), &written, 0);
+    return written;
   }
-  return written;
+
+  auto msg = utf8ToWChar(str);
+  return writeColorful(msg);
 }
 
 int WinConsoleFile::vprintf(const char* format, va_list va)
@@ -78,33 +132,134 @@ int WinConsoleFile::vprintf(const char* format, va_list va)
   if (r <= 0) {
     return 0;
   }
-  char *buf = new char[++r];
-  r = vsnprintf(buf, r, format, va);
+  auto buf = make_unique<char[]>(++r);
+  r = vsnprintf(buf.get(), r, format, va);
   if (r < 0) {
-    delete [] buf;
     return 0;
   }
-  DWORD written;
-  if(console(stdHandle_)) {
-    std::wstring msg = utf8ToWChar(buf);
-    WriteConsoleW(GetStdHandle(stdHandle_),
-                  msg.c_str(), msg.size(), &written, 0);
-  } else {
-    WriteFile(GetStdHandle(stdHandle_),
-              buf, r, &written, 0);
-  }
-  delete [] buf;
-  return written;
+  return write(buf.get());
 }
 
-int WinConsoleFile::flush()
+size_t WinConsoleFile::writeColorful(const std::wstring& str)
 {
-  return 0;
-}
+  size_t written = 0;
+  DWORD cw;
 
-bool WinConsoleFile::supportsColor()
-{
-  return false;
+  wchar_t suffix;
+  int arg = 0;
+  std::vector<int> args;
+  std::vector<wchar_t> buffer;
+  buffer.reserve(str.length());
+
+  enum state_ {
+    ePrefix, ePreFin, eNum0, eNum
+  } state = ePrefix;
+
+  for (const wchar_t ch : str) {
+    if (state == ePrefix) {
+      if (ch == '\033') {
+        state = ePreFin;
+      }
+      else {
+        buffer.push_back(ch);
+        continue;
+      }
+    }
+    else if (state == ePreFin) {
+      if (ch == '\033');
+      else if (ch == '[') {
+        state = eNum0;
+      }
+      else {
+        state = ePrefix;
+      }
+    }
+    else if (state == eNum0 || state == eNum) {
+      if (isdigit(ch)) {
+        arg = (arg * 10) + (ch - '0');
+        state = eNum;
+      }
+      else if (ch == ';') {
+        args.push_back(arg);
+        arg = 0;
+        state = eNum0;
+      }
+      else if (ch != '?') {
+        if (state == eNum) {
+          args.push_back(arg);
+        }
+        suffix = ch;
+        goto out;
+      }
+    }
+
+    ++written;
+    continue;
+
+out:
+    cw = 0;
+    if (!buffer.empty()) {
+      WriteConsoleW(handle(), buffer.data(), buffer.size(), &cw, nullptr);
+    }
+    written += cw;
+
+    if (suffix == 'm') {
+      if (args.empty()) {
+        args.push_back(0);
+      }
+      for (const int a: args) {
+        if (a == 0) {
+            fg_ = deffg_;
+            bg_ = defbg_;
+            bold_ = underline_ = reverse_ = false;
+        }
+        else if (30 <= a && a <= 37) {
+          fg_ = a - 30;
+        }
+        else if (40 <= a && a <= 47) {
+          bg_ = a - 40;
+        }
+        else if (a == 1 || a == 21) {
+          bold_ = a == 1;
+        }
+        else if (a == 4 || a == 24) {
+          underline_ = a == 4;
+        }
+        else if (a == 7 || a == 27) {
+          reverse_ = a == 7;
+        }
+      }
+      WORD attribute = 0;
+      if (reverse_) {
+        attribute = kForeground[bg_] | kBackground[fg_];
+      }
+      else {
+        attribute = kForeground[fg_] | kBackground[bg_];
+      }
+      if (bold_) {
+        attribute |= FOREGROUND_INTENSITY;
+      }
+      if (underline_) {
+        attribute |= BACKGROUND_INTENSITY;
+      }
+      SetConsoleTextAttribute(handle(), attribute);
+    }
+
+    suffix = 0;
+    state = ePrefix;
+    arg = 0;
+    args.clear();
+    buffer.clear();
+    ++written;
+  }
+
+  if (!buffer.empty()) {
+    cw = 0;
+    WriteConsoleW(handle(), buffer.data(), buffer.size(), &cw, nullptr);
+    written += cw;
+  }
+
+  return written;
 }
 
 } // namespace aria2