Merge pull request #139 from segevfiner/win32-vt-processing

Yet another attempt at supporting Windows 10's ANSI/VT console

colorama/ansitowin32.py

@@ -4,7 +4,7 @@
 import os
 
 from .ansi import AnsiFore, AnsiBack, AnsiStyle, Style, BEL
-from .winterm import WinTerm, WinColor, WinStyle
+from .winterm import enable_vt_processing, WinTerm, WinColor, WinStyle
 from .win32 import windll, winapi_test
 
 
@@ -94,15 +94,22 @@ def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
         # (e.g. Cygwin Terminal). In this case it's up to the terminal
         # to support the ANSI codes.
         conversion_supported = on_windows and winapi_test()
+        try:
+            fd = wrapped.fileno()
+        except Exception:
+            fd = -1
+        system_has_native_ansi = not on_windows or enable_vt_processing(fd)
+        have_tty = not self.stream.closed and self.stream.isatty()
+        need_conversion = conversion_supported and not system_has_native_ansi
 
         # should we strip ANSI sequences from our output?
         if strip is None:
-            strip = conversion_supported or (not self.stream.closed and not self.stream.isatty())
+            strip = need_conversion or not have_tty
         self.strip = strip
 
         # should we should convert ANSI sequences into win32 calls?
         if convert is None:
-            convert = conversion_supported and not self.stream.closed and self.stream.isatty()
+            convert = need_conversion and have_tty
         self.convert = convert
 
         # dict of ansi codes to win32 functions and parameters

colorama/tests/ansitowin32_test.py

@@ -1,13 +1,19 @@
 # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
 from io import StringIO, TextIOWrapper
 from unittest import TestCase, main
+try:
+    from contextlib import ExitStack
+except ImportError:
+    # python 2
+    from contextlib2 import ExitStack
 
 try:
     from unittest.mock import MagicMock, Mock, patch
 except ImportError:
     from mock import MagicMock, Mock, patch
 
 from ..ansitowin32 import AnsiToWin32, StreamWrapper
+from ..win32 import ENABLE_VIRTUAL_TERMINAL_PROCESSING
 from .utils import osname
 
 
@@ -239,5 +245,50 @@ def test_osc_codes(self):
                 stream.write(code)
             self.assertEqual(winterm.set_title.call_count, 2)
 
+    def test_native_windows_ansi(self):
+        with ExitStack() as stack:
+            def p(a, b):
+                stack.enter_context(patch(a, b, create=True))
+            # Pretend to be on Windows
+            p("colorama.ansitowin32.os.name", "nt")
+            p("colorama.ansitowin32.winapi_test", lambda: True)
+            p("colorama.win32.winapi_test", lambda: True)
+            p("colorama.winterm.win32.windll", "non-None")
+            p("colorama.winterm.get_osfhandle", lambda _: 1234)
+
+            # Pretend that our mock stream has native ANSI support
+            p(
+                "colorama.winterm.win32.GetConsoleMode",
+                lambda _: ENABLE_VIRTUAL_TERMINAL_PROCESSING,
+            )
+            SetConsoleMode = Mock()
+            p("colorama.winterm.win32.SetConsoleMode", SetConsoleMode)
+
+            stdout = Mock()
+            stdout.closed = False
+            stdout.isatty.return_value = True
+            stdout.fileno.return_value = 1
+
+            # Our fake console says it has native vt support, so AnsiToWin32 should
+            # enable that support and do nothing else.
+            stream = AnsiToWin32(stdout)
+            SetConsoleMode.assert_called_with(1234, ENABLE_VIRTUAL_TERMINAL_PROCESSING)
+            self.assertFalse(stream.strip)
+            self.assertFalse(stream.convert)
+            self.assertFalse(stream.should_wrap())
+
+            # Now let's pretend we're on an old Windows console, that doesn't have
+            # native ANSI support.
+            p("colorama.winterm.win32.GetConsoleMode", lambda _: 0)
+            SetConsoleMode = Mock()
+            p("colorama.winterm.win32.SetConsoleMode", SetConsoleMode)
+
+            stream = AnsiToWin32(stdout)
+            SetConsoleMode.assert_called_with(1234, ENABLE_VIRTUAL_TERMINAL_PROCESSING)
+            self.assertTrue(stream.strip)
+            self.assertTrue(stream.convert)
+            self.assertTrue(stream.should_wrap())
+
+
 if __name__ == '__main__':
     main()

colorama/win32.py

@@ -4,6 +4,8 @@
 STDOUT = -11
 STDERR = -12
 
+ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
+
 try:
     import ctypes
     from ctypes import LibraryLoader
@@ -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
+
     def _winapi_test(handle):
         csbi = CONSOLE_SCREEN_BUFFER_INFO()
         success = _GetConsoleScreenBufferInfo(
@@ -150,3 +166,15 @@ def FillConsoleOutputAttribute(stream_id, attr, length, start):
 
     def SetConsoleTitle(title):
         return _SetConsoleTitleW(title)
+
+    def GetConsoleMode(handle):
+        mode = wintypes.DWORD()
+        success = _GetConsoleMode(handle, byref(mode))
+        if not success:
+            raise ctypes.WinError()
+        return mode.value
+
+    def SetConsoleMode(handle, mode):
+        success = _SetConsoleMode(handle, mode)
+        if not success:
+            raise ctypes.WinError()

colorama/winterm.py

@@ -1,6 +1,12 @@
 # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
-from . import win32
+try:
+    from msvcrt import get_osfhandle
+except ImportError:
+    def get_osfhandle(_):
+        raise OSError("This isn't windows!")
+
 
+from . import win32
 
 # from wincon.h
 class WinColor(object):
@@ -167,3 +173,22 @@ def erase_line(self, mode=0, on_stderr=False):
 
     def set_title(self, title):
         win32.SetConsoleTitle(title)
+
+
+def enable_vt_processing(fd):
+    if win32.windll is None or not win32.winapi_test():
+        return False
+
+    try:
+        handle = get_osfhandle(fd)
+        mode = win32.GetConsoleMode(handle)
+        win32.SetConsoleMode(
+            handle,
+            mode | win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING,
+        )
+
+        mode = win32.GetConsoleMode(handle)
+        if mode & win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING:
+            return True
+    except OSError:
+        return False

requirements-dev.txt

@@ -1,4 +1,5 @@
 mock>=1.0.1;python_version<"3.3"
 twine>=3.1.1
+contextlib2;python_version<"3"
 build
 -e .

tox.ini

@@ -3,5 +3,7 @@ isolated_build = true
 envlist = py27, py37, py38, py39, py310, pypy, pypy3
 
 [testenv]
-deps = py27,pypy: mock
+deps =
+  py27,pypy: mock
+  py27,pypy: contextlib2
 commands = python -m unittest discover -p *_test.py