diff --git a/colorama/ansitowin32.py b/colorama/ansitowin32.py index 6039a05..9941f6d 100644 --- a/colorama/ansitowin32.py +++ b/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 @@ -80,6 +80,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 @@ -89,12 +91,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 self.stream.closed and not self.stream.isatty()) + strip = ((not vt_processing_supported and conversion_supported) or + (not self.stream.closed and not self.stream.isatty())) 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 = (not vt_processing_supported and conversion_supported and + not self.stream.closed and self.stream.isatty()) self.convert = convert # dict of ansi codes to win32 functions and parameters diff --git a/colorama/win32.py b/colorama/win32.py index c2d8360..1aeb639 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 + def _winapi_test(handle): csbi = CONSOLE_SCREEN_BUFFER_INFO() success = _GetConsoleScreenBufferInfo( @@ -150,3 +166,13 @@ def FillConsoleOutputAttribute(stream_id, attr, length, start): def SetConsoleTitle(title): return _SetConsoleTitleW(title) + + def GetConsoleMode(stream_id): + handle = _GetStdHandle(stream_id) + mode = wintypes.DWORD() + success = _GetConsoleMode(handle, byref(mode)) + return mode.value + + def SetConsoleMode(stream_id, mode): + handle = _GetStdHandle(stream_id) + return _SetConsoleMode(handle, mode) diff --git a/colorama/winterm.py b/colorama/winterm.py index 0fdb4ec..95b36c9 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 @@ -167,3 +168,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