Use ENABLE_VIRTUAL_TERMINAL_PROCESSING when available

colorama/ansitowin32.py

@@ -4,7 +4,7 @@
 import os
 
 from .ansi import AnsiFore, AnsiBack, AnsiStyle, Style
-from .winterm import WinTerm, WinColor, WinStyle
+from .winterm import enable_vt_processing, WinTerm, WinColor, WinStyle
 from .win32 import windll, winapi_test
 
 
@@ -59,6 +59,8 @@ def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
         # create the proxy wrapping our output stream
         self.stream = StreamWrapper(wrapped, self)
 
+        vt_processing_supported = enable_vt_processing()
+
         on_windows = os.name == 'nt'
         # We test if the WinAPI works, because even if we are on Windows
         # we may be using a terminal that doesn't support the WinAPI
@@ -68,12 +70,14 @@ def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
 
         # should we strip ANSI sequences from our output?
         if strip is None:
-            strip = conversion_supported or (not is_stream_closed(wrapped) and not is_a_tty(wrapped))
+            strip = ((not vt_processing_supported and conversion_supported) or
+                     (not is_stream_closed(wrapped) and not is_a_tty(wrapped)))
         self.strip = strip
 
         # should we should convert ANSI sequences into win32 calls?
         if convert is None:
-            convert = conversion_supported and not is_stream_closed(wrapped) and is_a_tty(wrapped)
+            convert = (not vt_processing_supported and conversion_supported and
+                       not is_stream_closed(wrapped) and is_a_tty(wrapped))
         self.convert = convert
 
         # dict of ansi codes to win32 functions and parameters

colorama/win32.py

@@ -16,6 +16,8 @@
 else:
     from ctypes import byref, Structure, c_char, POINTER
 
+    ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
+
     COORD = wintypes._COORD
 
     class CONSOLE_SCREEN_BUFFER_INFO(Structure):
@@ -89,6 +91,20 @@ def __str__(self):
     ]
     _SetConsoleTitleW.restype = wintypes.BOOL
 
+    _GetConsoleMode = windll.kernel32.GetConsoleMode
+    _GetConsoleMode.argtypes = [
+        wintypes.HANDLE,
+        POINTER(wintypes.DWORD)
+    ]
+    _GetConsoleMode.restype = wintypes.BOOL
+
+    _SetConsoleMode = windll.kernel32.SetConsoleMode
+    _SetConsoleMode.argtypes = [
+        wintypes.HANDLE,
+        wintypes.DWORD
+    ]
+    _SetConsoleMode.restype = wintypes.BOOL
+
     handles = {
         STDOUT: _GetStdHandle(STDOUT),
         STDERR: _GetStdHandle(STDERR),
@@ -154,3 +170,13 @@ def FillConsoleOutputAttribute(stream_id, attr, length, start):
 
     def SetConsoleTitle(title):
         return _SetConsoleTitleW(title)
+
+    def GetConsoleMode(stream_id):
+        handle = handles[stream_id]
+        mode = wintypes.DWORD()
+        success = _GetConsoleMode(handle, byref(mode))
+        return mode.value
+
+    def SetConsoleMode(stream_id, mode):
+        handle = handles[stream_id]
+        return _SetConsoleMode(handle, mode)

colorama/winterm.py

@@ -1,4 +1,5 @@
 # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
+import atexit
 from . import win32
 
 
@@ -160,3 +161,48 @@ def erase_line(self, mode=0, on_stderr=False):
 
     def set_title(self, title):
         win32.SetConsoleTitle(title)
+
+
+_enabled_vt_processing = False
+_atexit_registered = False
+
+
+def enable_vt_processing():
+    global _enabled_vt_processing, _atexit_registered
+
+    if win32.windll is None or not win32.winapi_test():
+        return False
+
+    mode = win32.GetConsoleMode(win32.STDOUT)
+
+    if mode & win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING:
+        return True
+
+    if not win32.SetConsoleMode(
+        win32.STDOUT,
+        mode | win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING):
+            return False  # Unsupported
+
+    _enabled_vt_processing = True
+
+    if not _atexit_registered:
+        atexit.register(lambda: restore_vt_processing(True))
+        _atexit_registered = True
+
+    return True
+
+
+def restore_vt_processing(atexit=False):
+    global _enabled_vt_processing, _atexit_registered
+
+    if atexit:
+        _atexit_registered = False
+
+    if win32.windll is None or not win32.winapi_test():
+        return
+
+    if _enabled_vt_processing:
+        win32.SetConsoleMode(
+            win32.STDOUT,
+            win32.GetConsoleMode(win32.STDOUT) & ~win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING)
+        _enabled_vt_processing = False