tools: make color output work on windows

6 files changed, 115 insertions, 4 deletions

diff --git a/tools/build.py b/tools/build.py
index 9de84f541..bbcd24762 100755
--- a/tools/build.py
+++ b/tools/build.py
@@ -4,7 +4,9 @@ from __future__ import print_function
 import os
 import sys
 import third_party
-from util import build_path, run
+from util import build_path, enable_ansi_colors, run
+
+enable_ansi_colors()
 
 third_party.fix_symlinks()
 
diff --git a/tools/lint.py b/tools/lint.py
index 0ed422241..e5a4f1690 100755
--- a/tools/lint.py
+++ b/tools/lint.py
@@ -2,7 +2,9 @@
 # Does google-lint on c++ files and ts-lint on typescript files
 
 import os
-from util import run
+from util import enable_ansi_colors, run
+
+enable_ansi_colors()
 
 root_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
 third_party_path = os.path.join(root_path, "third_party")
diff --git a/tools/setup.py b/tools/setup.py
index 77aaf95d6..1f065769a 100755
--- a/tools/setup.py
+++ b/tools/setup.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 import third_party
-from util import run, root_path, build_path, build_mode
+from util import build_mode, build_path, enable_ansi_colors, root_path, run
 import os
 import re
 import sys
@@ -8,6 +8,8 @@ from distutils.spawn import find_executable
 
 
 def main():
+    enable_ansi_colors()
+
     os.chdir(root_path)
 
     third_party.fix_symlinks()
diff --git a/tools/sync_third_party.py b/tools/sync_third_party.py
index cf0de30f3..697470a49 100755
--- a/tools/sync_third_party.py
+++ b/tools/sync_third_party.py
@@ -5,6 +5,9 @@
 # find . -type f | grep -v "\.git" | xargs -I% git add -f --no-warn-embedded-repo "%"
 
 import third_party
+import util
+
+util.enable_ansi_colors()
 
 third_party.fix_symlinks()
 
diff --git a/tools/test.py b/tools/test.py
index b8f49cf2b..3cd03e916 100755
--- a/tools/test.py
+++ b/tools/test.py
@@ -5,7 +5,7 @@ import os
 import sys
 from check_output_test import check_output_test
 from setup_test import setup_test
-from util import executable_suffix, run, build_path
+from util import build_path, enable_ansi_colors, executable_suffix, run
 from unit_tests import unit_tests
 from util_test import util_test
 import subprocess
@@ -28,6 +28,8 @@ def main(argv):
         print "Usage: tools/test.py [build_dir]"
         sys.exit(1)
 
+    enable_ansi_colors()
+
     http_server.spawn()
 
     # Internal tools testing
diff --git a/tools/util.py b/tools/util.py
index b50b754c2..659d9584b 100644
--- a/tools/util.py
+++ b/tools/util.py
@@ -173,3 +173,103 @@ def parse_exit_code(s):
         return codes[0]
     else:
         return 0
+
+
+# Attempts to enable ANSI escape code support.
+# Returns True if successful, False if not supported.
+def enable_ansi_colors():
+    if os.name != 'nt':
+        return True  # On non-windows platforms this just works.
+    elif "CI" in os.environ:
+        return True  # Ansi escape codes work out of the box on Appveyor.
+
+    return enable_ansi_colors_win10()
+
+
+# The windows 10 implementation of enable_ansi_colors.
+def enable_ansi_colors_win10():
+    import ctypes
+
+    # Function factory for errcheck callbacks that raise WinError on failure.
+    def raise_if(error_result):
+        def check(result, func, args):
+            if result == error_result:
+                raise ctypes.WinError(ctypes.get_last_error())
+            return args
+
+        return check
+
+    # Windows API types.
+    from ctypes.wintypes import BOOL, DWORD, HANDLE, LPCWSTR, LPVOID
+    LPDWORD = ctypes.POINTER(DWORD)
+
+    # Generic constants.
+    NULL = ctypes.c_void_p(0).value
+    INVALID_HANDLE_VALUE = ctypes.c_void_p(-1).value
+
+    # CreateFile flags.
+    # yapf: disable
+    GENERIC_READ  = 0x80000000
+    GENERIC_WRITE = 0x40000000
+    FILE_SHARE_READ  = 0x01
+    FILE_SHARE_WRITE = 0x02
+    OPEN_EXISTING = 3
+    # yapf: enable
+
+    # Get/SetConsoleMode flags.
+    ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x04
+
+    kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
+
+    # HANDLE CreateFileW(...)
+    CreateFileW = kernel32.CreateFileW
+    CreateFileW.restype = HANDLE
+    CreateFileW.errcheck = raise_if(INVALID_HANDLE_VALUE)
+    # yapf: disable
+    CreateFileW.argtypes = (LPCWSTR,  # lpFileName
+                            DWORD,    # dwDesiredAccess
+                            DWORD,    # dwShareMode
+                            LPVOID,   # lpSecurityAttributes
+                            DWORD,    # dwCreationDisposition
+                            DWORD,    # dwFlagsAndAttributes
+                            HANDLE)   # hTemplateFile
+    # yapf: enable
+
+    # BOOL CloseHandle(HANDLE hObject)
+    CloseHandle = kernel32.CloseHandle
+    CloseHandle.restype = BOOL
+    CloseHandle.errcheck = raise_if(False)
+    CloseHandle.argtypes = (HANDLE, )
+
+    # BOOL GetConsoleMode(HANDLE hConsoleHandle, LPDWORD lpMode)
+    GetConsoleMode = kernel32.GetConsoleMode
+    GetConsoleMode.restype = BOOL
+    GetConsoleMode.errcheck = raise_if(False)
+    GetConsoleMode.argtypes = (HANDLE, LPDWORD)
+
+    # BOOL SetConsoleMode(HANDLE hConsoleHandle, DWORD dwMode)
+    SetConsoleMode = kernel32.SetConsoleMode
+    SetConsoleMode.restype = BOOL
+    SetConsoleMode.errcheck = raise_if(False)
+    SetConsoleMode.argtypes = (HANDLE, DWORD)
+
+    # Open the console output device.
+    conout = CreateFileW("CONOUT$", GENERIC_READ | GENERIC_WRITE,
+                         FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
+                         OPEN_EXISTING, 0, 0)
+
+    # Get the current mode.
+    mode = DWORD()
+    GetConsoleMode(conout, ctypes.byref(mode))
+
+    # Try to set the flag that controls ANSI escape code support.
+    try:
+        SetConsoleMode(conout, mode.value | ENABLE_VIRTUAL_TERMINAL_PROCESSING)
+    except WindowsError as e:
+        if e.winerror == ERROR_INVALID_PARAMETER:
+            return False  # Not supported, likely an older version of Windows.
+        raise
+    finally:
+        CloseHandle(conout)
+
+    return True