Merge pull request #277 from zmwangx/smart-color

Smarter colorization and better support for native terminals on Windows

googler

@@ -30,6 +30,7 @@ from http.client import HTTPSConnection
 import locale
 import logging
 import os
+import platform
 import shutil
 import signal
 import socket
@@ -3126,8 +3127,13 @@ def parse_args(args=None, namespace=None):
     addarg('-l', '--lang', metavar='LANG', help='display in language LANG')
     addarg('-x', '--exact', action='store_true',
            help='disable automatic spelling correction')
-    addarg('-C', '--nocolor', dest='colorize', action='store_false',
-           help='disable color output')
+    addarg('--colorize', nargs='?', choices=['auto', 'always', 'never'],
+           const='always', default='auto',
+           help="""whether to colorize output; defaults to 'auto', which enables
+           color when stdout is a tty device; using --colorize without an argument
+           is equivalent to --colorize=always""")
+    addarg('-C', '--nocolor', action='store_true',
+           help='equivalent to --colorize=never')
     addarg('--colors', dest='colorstr', type=argparser.is_colorstr,
            default=colorstr_env if colorstr_env else 'GKlgxy', metavar='COLORS',
            help='set output colors (see man page for details)')
@@ -3165,7 +3171,11 @@ def parse_args(args=None, namespace=None):
     addarg('-D', '--debugger', action='store_true', help=argparse.SUPPRESS)
     addarg('--complete', help=argparse.SUPPRESS)
 
-    return argparser.parse_args(args, namespace)
+    parsed = argparser.parse_args(args, namespace)
+    if parsed.nocolor:
+        parsed.colorize = 'never'
+
+    return parsed
 
 
 def main():
@@ -3203,14 +3213,46 @@ def main():
                 pass
 
         # Set colors
-        if opts.colorize:
+        if opts.colorize == 'always':
+            colorize = True
+        elif opts.colorize == 'auto':
+            colorize = sys.stdout.isatty()
+        elif opts.colorize == 'never':
+            colorize = False
+        else:
+            raise ValueError("invalid --colorize value '%s'" % opts.colorize)
+
+        if colorize:
             colors = Colors(*[COLORMAP[c] for c in opts.colorstr], reset=COLORMAP['x'])
         else:
             colors = None
         Result.colors = colors
         Result.urlexpand = True if os.getenv('DISABLE_URL_EXPANSION') is None else False
         GooglerCmd.colors = colors
 
+        # Try to enable ANSI color support in cmd or PowerShell on Windows 10
+        if sys.platform == 'win32' and sys.stdout.isatty() and colorize:
+            # VT100 control sequences are supported on Windows 10 Anniversary Update and later.
+            # https://docs.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences
+            # https://docs.microsoft.com/en-us/windows/console/setconsolemode
+            if platform.release() == '10':
+                STD_OUTPUT_HANDLE = -11
+                ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
+                try:
+                    from ctypes import windll, wintypes, byref
+                    kernel32 = windll.kernel32
+                    stdout_handle = kernel32.GetStdHandle(STD_OUTPUT_HANDLE)
+                    old_mode = wintypes.DWORD()
+                    if not kernel32.GetConsoleMode(stdout_handle, byref(old_mode)):
+                        raise RuntimeError('GetConsoleMode failed')
+                    new_mode = old_mode.value | ENABLE_VIRTUAL_TERMINAL_PROCESSING
+                    if not kernel32.SetConsoleMode(stdout_handle, new_mode):
+                        raise RuntimeError('SetConsoleMode failed')
+                    # Note: No need to restore at exit. SetConsoleMode seems to
+                    # be limited to the calling process.
+                except Exception:
+                    pass
+
         if opts.url_handler is not None:
             open_url.url_handler = opts.url_handler
         else: