From 6e774c68313e85c9f8a46bdbf6fa48be502d82f0 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Wed, 4 Dec 2019 18:27:55 +0530 Subject: [PATCH] Simplify windows terminal handling code We rely on the fact that python 3 now knows how to write unicode to windows consoles. And we setup the console to understand ANSI escape codes --- src/calibre/utils/terminal.py | 209 ++++++---------------------------- 1 file changed, 32 insertions(+), 177 deletions(-) diff --git a/src/calibre/utils/terminal.py b/src/calibre/utils/terminal.py index c61f45fad9..51d38580bd 100644 --- a/src/calibre/utils/terminal.py +++ b/src/calibre/utils/terminal.py @@ -5,10 +5,10 @@ __license__ = 'GPL v3' __copyright__ = '2012, Kovid Goyal ' __docformat__ = 'restructuredtext en' -import os, sys, re +import os, sys, re, io from calibre.constants import iswindows -from polyglot.builtins import iteritems, range, zip, native_string_type +from polyglot.builtins import iteritems, range, zip if iswindows: import ctypes.wintypes @@ -70,21 +70,6 @@ COLORS = {v:fmt(k) for k, v in iteritems(RCOLORS)} RESET = fmt(0) -if iswindows: - # From wincon.h - WCOLORS = {c:i for i, c in enumerate(( - 'black', 'blue', 'green', 'cyan', 'red', 'magenta', 'yellow', 'white'))} - - def to_flag(fg, bg, bold): - val = 0 - if bold: - val |= 0x08 - if fg in WCOLORS: - val |= WCOLORS[fg] - if bg in WCOLORS: - val |= (WCOLORS[bg] << 4) - return val - def colored(text, fg=None, bg=None, bold=False): prefix = [] @@ -102,93 +87,38 @@ def colored(text, fg=None, bg=None, bold=False): return prefix + text + suffix +def is_binary(stream): + mode = getattr(stream, 'mode', None) + if mode: + return 'b' in mode + return not isinstance(stream, io.TextIOBase) + + class Detect(object): def __init__(self, stream): self.stream = stream or sys.stdout + self.is_binary = is_binary(self.stream) self.isatty = getattr(self.stream, 'isatty', lambda : False)() force_ansi = 'CALIBRE_FORCE_ANSI' in os.environ if not self.isatty and force_ansi: self.isatty = True - self.isansi = force_ansi or not iswindows - self.set_console = self.write_console = None - self.is_console = False - if not self.isansi: - try: - import msvcrt - self.msvcrt = msvcrt - self.file_handle = msvcrt.get_osfhandle(self.stream.fileno()) - from ctypes import windll, wintypes, byref, POINTER, WinDLL, c_wchar_p - mode = wintypes.DWORD(0) - f = windll.kernel32.GetConsoleMode - f.argtypes, f.restype = [wintypes.HANDLE, POINTER(wintypes.DWORD)], wintypes.BOOL - if f(self.file_handle, byref(mode)): - # Stream is a console - self.set_console = windll.kernel32.SetConsoleTextAttribute - self.default_console_text_attributes = WCOLORS['white'] - kernel32 = WinDLL(native_string_type('kernel32'), use_last_error=True) - self.write_console = kernel32.WriteConsoleW - self.write_console.argtypes = [wintypes.HANDLE, c_wchar_p, wintypes.DWORD, POINTER(wintypes.DWORD), wintypes.LPVOID] - self.write_console.restype = wintypes.BOOL - kernel32.GetConsoleScreenBufferInfo.argtypes = [wintypes.HANDLE, ctypes.POINTER(CONSOLE_SCREEN_BUFFER_INFO)] - kernel32.GetConsoleScreenBufferInfo.restype = wintypes.BOOL - csbi = CONSOLE_SCREEN_BUFFER_INFO() - if kernel32.GetConsoleScreenBufferInfo(self.file_handle, byref(csbi)): - self.default_console_text_attributes = csbi.wAttributes - self.is_console = True - except: - pass - - def write_unicode_text(self, text, ignore_errors=False): - ' Windows only method that writes unicode strings correctly to the windows console using the Win32 API ' - if self.is_console: - from ctypes import wintypes, byref, c_wchar_p - written = wintypes.DWORD(0) - text = text.replace('\0', '') - chunk = len(text) - while text: - t, text = text[:chunk], text[chunk:] - wt = c_wchar_p(t) - text_len = len(t.encode('utf-16')) - if not self.write_console(self.file_handle, wt, text_len, byref(written), None): - # Older versions of windows can fail to write large strings - # to console with WriteConsoleW (seen it happen on Win XP) - import winerror - err = ctypes.get_last_error() - if err == winerror.ERROR_NOT_ENOUGH_MEMORY and chunk >= 128: - # Retry with a smaller chunk size (give up if chunk < 128) - chunk = chunk // 2 - text = t + text - continue - if err == winerror.ERROR_GEN_FAILURE: - # On newer windows, this happens when trying to write - # non-ascii chars to the console and the console is set - # to use raster fonts (the default). In this case - # rather than failing, write an informative error - # message and the asciized version of the text. - print('Non-ASCII text detected. You must set your Console\'s font to' - ' Lucida Console or Consolas or some other TrueType font to see this text', file=self.stream, end=' -- ') - from calibre.utils.filenames import ascii_text - print(ascii_text(t + text), file=self.stream, end='') - continue - if not ignore_errors: - raise ctypes.WinError(err) + self.isansi = force_ansi or not iswindows or (iswindows and sys.getwindowsversion().major >= 10) class ColoredStream(Detect): def __init__(self, stream=None, fg=None, bg=None, bold=False): - stream = getattr(stream, 'buffer', stream) Detect.__init__(self, stream) self.fg, self.bg, self.bold = fg, bg, bold - if self.set_console is not None: - self.wval = to_flag(self.fg, self.bg, bold) - if not self.bg: - self.wval |= self.default_console_text_attributes & 0xF0 def cwrite(self, what): - if not isinstance(what, bytes): - what = what.encode('ascii') + if self.is_binary: + if not isinstance(what, bytes): + what = what.encode('utf-8') + else: + if isinstance(what, bytes): + what = what.decode('utf-8', 'replace') self.stream.write(what) def __enter__(self): @@ -201,9 +131,6 @@ class ColoredStream(Detect): self.cwrite(BACKGROUNDS[self.bg]) if self.fg is not None: self.cwrite(COLORS[self.fg]) - elif self.set_console is not None: - if self.wval != 0: - self.set_console(self.file_handle, self.wval) return self def __exit__(self, *args, **kwargs): @@ -214,8 +141,6 @@ class ColoredStream(Detect): if self.isansi: self.cwrite(RESET) self.stream.flush() - elif self.set_console is not None: - self.set_console(self.file_handle, self.default_console_text_attributes) class ANSIStream(Detect): @@ -224,9 +149,7 @@ class ANSIStream(Detect): def __init__(self, stream=None): super(ANSIStream, self).__init__(stream) - self.encoding = getattr(self.stream, 'encoding', 'utf-8') or 'utf-8' - self.stream_takes_unicode = hasattr(self.stream, 'buffer') - self.last_state = (None, None, False) + self.encoding = getattr(self.stream, 'encoding', None) or 'utf-8' self._ansi_re_bin = self._ansi_re_unicode = None def ansi_re(self, binary=False): @@ -245,97 +168,29 @@ class ANSIStream(Detect): return self.strip_and_write(text) if self.isansi: - return self.stream.write(text) + return self.polyglot_write(text) - if not self.isansi and self.set_console is None: - return self.strip_and_write(text) - - self.write_and_convert(text) + return self.strip_and_write(text) def polyglot_write(self, text): binary = isinstance(text, bytes) stream = self.stream - if self.stream_takes_unicode: - if binary: - stream = self.stream.buffer - else: - if not binary: - text = text.encode(self.encoding, 'replace') - stream.write(text) + if binary: + if self.is_binary: + return stream.write(text) + buffer = getattr(stream, 'buffer', None) + if buffer is None: + return stream.write(text.decode('utf-8', 'replace')) + return buffer.write(text) + if self.is_binary: + text = text.encode(self.encoding, 'replace') + return stream.write(text) def strip_and_write(self, text): binary = isinstance(text, bytes) pat = self.ansi_re(binary) repl = b'' if binary else '' - self.polyglot_write(pat.sub(repl, text)) - - def write_and_convert(self, text): - ''' - Write the given text to our wrapped stream, stripping any ANSI - sequences from the text, and optionally converting them into win32 - calls. - ''' - cursor = 0 - binary = isinstance(text, bytes) - for match in self.ansi_re(binary).finditer(text): - start, end = match.span() - self.write_plain_text(text, cursor, start) - self.convert_ansi(*match.groups()) - cursor = end - self.write_plain_text(text, cursor, len(text)) - self.set_console(self.file_handle, self.default_console_text_attributes) - self.stream.flush() - - def write_plain_text(self, text, start, end): - if start < end: - text = text[start:end] - if self.is_console and isinstance(text, bytes): - try: - utext = text.decode(self.encoding) - except ValueError: - pass - else: - return self.write_unicode_text(utext) - self.polyglot_write(text) - - def convert_ansi(self, paramstring, command): - if isinstance(paramstring, bytes): - paramstring = paramstring.decode('ascii', 'replace') - if isinstance(command, bytes): - command = command.decode('ascii', 'replace') - params = self.extract_params(paramstring) - self.call_win32(command, params) - - def extract_params(self, paramstring): - def split(paramstring): - for p in paramstring.split(';'): - if p: - yield int(p) - return tuple(split(paramstring)) - - def call_win32(self, command, params): - if command != 'm': - return - fg, bg, bold = self.last_state - - for param in params: - if param in RCOLORS: - fg = RCOLORS[param] - elif param in RBACKGROUNDS: - bg = RBACKGROUNDS[param] - elif param == 1: - bold = True - elif param == 0: - fg, bg, bold = None, None, False - - self.last_state = (fg, bg, bold) - if fg or bg or bold: - val = to_flag(fg, bg, bold) - if not bg: - val |= self.default_console_text_attributes & 0xF0 - self.set_console(self.file_handle, val) - else: - self.set_console(self.file_handle, self.default_console_text_attributes) + return self.polyglot_write(pat.sub(repl, text)) def windows_terminfo(): @@ -433,5 +288,5 @@ def test(): s.write('\n'.join(text)) u = '\u041c\u0438\u0445\u0430\u0438\u043b fällen' print() - s.write_unicode_text(u) + s.write(u) print()