Fix native ANSI when stderr=console, stdout=not console + basic test

colorama/ansitowin32.py

@@ -86,9 +86,13 @@ 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()
-        native_ansi_supported = not on_windows or enable_vt_processing()
+        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 native_ansi_supported
+        need_conversion = conversion_supported and not system_has_native_ansi
 
         # should we strip ANSI sequences from our output?
         if strip is None:

colorama/tests/ansitowin32_test.py

@@ -1,13 +1,15 @@
 # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
 from io import StringIO
 from unittest import TestCase, main
+from contextlib 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
 
 
@@ -228,5 +230,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
@@ -16,8 +18,6 @@
 else:
     from ctypes import byref, Structure, c_char, POINTER
 
-    ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
-
     COORD = wintypes._COORD
 
     class CONSOLE_SCREEN_BUFFER_INFO(Structure):
@@ -167,16 +167,14 @@ def FillConsoleOutputAttribute(stream_id, attr, length, start):
     def SetConsoleTitle(title):
         return _SetConsoleTitleW(title)
 
-    def GetConsoleMode(stream_id):
-        handle = _GetStdHandle(stream_id)
+    def GetConsoleMode(handle):
         mode = wintypes.DWORD()
         success = _GetConsoleMode(handle, byref(mode))
         if not success:
             raise ctypes.WinError()
         return mode.value
 
-    def SetConsoleMode(stream_id, mode):
-        handle = _GetStdHandle(stream_id)
+    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):
@@ -169,21 +175,19 @@ def set_title(self, title):
         win32.SetConsoleTitle(title)
 
 
-def enable_vt_processing():
+def enable_vt_processing(fd):
     if win32.windll is None or not win32.winapi_test():
         return False
 
     try:
-        mode = win32.GetConsoleMode(win32.STDOUT)
-        if mode & win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING:
-            return True
-
+        handle = get_osfhandle(fd)
+        mode = win32.GetConsoleMode(handle)
         win32.SetConsoleMode(
-            win32.STDOUT,
+            handle,
             mode | win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING,
         )
 
-        mode = win32.GetConsoleMode(win32.STDOUT)
+        mode = win32.GetConsoleMode(handle)
         if mode & win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING:
             return True
     except OSError: