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
This commit is contained in:
Kovid Goyal 2019-12-04 18:27:55 +05:30
parent df70399661
commit 6e774c6831
No known key found for this signature in database
GPG Key ID: 06BC317B515ACE7C

View File

@ -5,10 +5,10 @@ __license__ = 'GPL v3'
__copyright__ = '2012, Kovid Goyal <kovid at kovidgoyal.net>'
__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()