diff --git a/colorama/ansitowin32.py b/colorama/ansitowin32.py index 1d6e605..68e5db7 100644 --- a/colorama/ansitowin32.py +++ b/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 diff --git a/colorama/win32.py b/colorama/win32.py index 8262e35..7aa3d91 100644 --- a/colorama/win32.py +++ b/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) diff --git a/colorama/winterm.py b/colorama/winterm.py index 60309d3..30caad3 100644 --- a/colorama/winterm.py +++ b/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