diff --git a/colorama/ansitowin32.py b/colorama/ansitowin32.py index 2b7b2c1..0cf7c39 100644 --- a/colorama/ansitowin32.py +++ b/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: diff --git a/colorama/tests/ansitowin32_test.py b/colorama/tests/ansitowin32_test.py index bbe99f4..66a2266 100644 --- a/colorama/tests/ansitowin32_test.py +++ b/colorama/tests/ansitowin32_test.py @@ -1,6 +1,7 @@ # 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 @@ -8,6 +9,7 @@ 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() diff --git a/colorama/win32.py b/colorama/win32.py index 6bd96b9..841b0e2 100644 --- a/colorama/win32.py +++ b/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() diff --git a/colorama/winterm.py b/colorama/winterm.py index 0308cb0..fd7202c 100644 --- a/colorama/winterm.py +++ b/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: