123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423 |
- """
- Helper functions for writing to terminals and files.
- """
- import sys, os, unicodedata
- import py
- py3k = sys.version_info[0] >= 3
- py33 = sys.version_info >= (3, 3)
- from py.builtin import text, bytes
- win32_and_ctypes = False
- colorama = None
- if sys.platform == "win32":
- try:
- import colorama
- except ImportError:
- try:
- import ctypes
- win32_and_ctypes = True
- except ImportError:
- pass
- def _getdimensions():
- if py33:
- import shutil
- size = shutil.get_terminal_size()
- return size.lines, size.columns
- else:
- import termios, fcntl, struct
- call = fcntl.ioctl(1, termios.TIOCGWINSZ, "\000" * 8)
- height, width = struct.unpack("hhhh", call)[:2]
- return height, width
- def get_terminal_width():
- width = 0
- try:
- _, width = _getdimensions()
- except py.builtin._sysex:
- raise
- except:
- # pass to fallback below
- pass
- if width == 0:
- # FALLBACK:
- # * some exception happened
- # * or this is emacs terminal which reports (0,0)
- width = int(os.environ.get('COLUMNS', 80))
- # XXX the windows getdimensions may be bogus, let's sanify a bit
- if width < 40:
- width = 80
- return width
- terminal_width = get_terminal_width()
- char_width = {
- 'A': 1, # "Ambiguous"
- 'F': 2, # Fullwidth
- 'H': 1, # Halfwidth
- 'N': 1, # Neutral
- 'Na': 1, # Narrow
- 'W': 2, # Wide
- }
- def get_line_width(text):
- text = unicodedata.normalize('NFC', text)
- return sum(char_width.get(unicodedata.east_asian_width(c), 1) for c in text)
- # XXX unify with _escaped func below
- def ansi_print(text, esc, file=None, newline=True, flush=False):
- if file is None:
- file = sys.stderr
- text = text.rstrip()
- if esc and not isinstance(esc, tuple):
- esc = (esc,)
- if esc and sys.platform != "win32" and file.isatty():
- text = (''.join(['\x1b[%sm' % cod for cod in esc]) +
- text +
- '\x1b[0m') # ANSI color code "reset"
- if newline:
- text += '\n'
- if esc and win32_and_ctypes and file.isatty():
- if 1 in esc:
- bold = True
- esc = tuple([x for x in esc if x != 1])
- else:
- bold = False
- esctable = {() : FOREGROUND_WHITE, # normal
- (31,): FOREGROUND_RED, # red
- (32,): FOREGROUND_GREEN, # green
- (33,): FOREGROUND_GREEN|FOREGROUND_RED, # yellow
- (34,): FOREGROUND_BLUE, # blue
- (35,): FOREGROUND_BLUE|FOREGROUND_RED, # purple
- (36,): FOREGROUND_BLUE|FOREGROUND_GREEN, # cyan
- (37,): FOREGROUND_WHITE, # white
- (39,): FOREGROUND_WHITE, # reset
- }
- attr = esctable.get(esc, FOREGROUND_WHITE)
- if bold:
- attr |= FOREGROUND_INTENSITY
- STD_OUTPUT_HANDLE = -11
- STD_ERROR_HANDLE = -12
- if file is sys.stderr:
- handle = GetStdHandle(STD_ERROR_HANDLE)
- else:
- handle = GetStdHandle(STD_OUTPUT_HANDLE)
- oldcolors = GetConsoleInfo(handle).wAttributes
- attr |= (oldcolors & 0x0f0)
- SetConsoleTextAttribute(handle, attr)
- while len(text) > 32768:
- file.write(text[:32768])
- text = text[32768:]
- if text:
- file.write(text)
- SetConsoleTextAttribute(handle, oldcolors)
- else:
- file.write(text)
- if flush:
- file.flush()
- def should_do_markup(file):
- if os.environ.get('PY_COLORS') == '1':
- return True
- if os.environ.get('PY_COLORS') == '0':
- return False
- if 'NO_COLOR' in os.environ:
- return False
- return hasattr(file, 'isatty') and file.isatty() \
- and os.environ.get('TERM') != 'dumb' \
- and not (sys.platform.startswith('java') and os._name == 'nt')
- class TerminalWriter(object):
- _esctable = dict(black=30, red=31, green=32, yellow=33,
- blue=34, purple=35, cyan=36, white=37,
- Black=40, Red=41, Green=42, Yellow=43,
- Blue=44, Purple=45, Cyan=46, White=47,
- bold=1, light=2, blink=5, invert=7)
- # XXX deprecate stringio argument
- def __init__(self, file=None, stringio=False, encoding=None):
- if file is None:
- if stringio:
- self.stringio = file = py.io.TextIO()
- else:
- from sys import stdout as file
- elif py.builtin.callable(file) and not (
- hasattr(file, "write") and hasattr(file, "flush")):
- file = WriteFile(file, encoding=encoding)
- if hasattr(file, "isatty") and file.isatty() and colorama:
- file = colorama.AnsiToWin32(file).stream
- self.encoding = encoding or getattr(file, 'encoding', "utf-8")
- self._file = file
- self.hasmarkup = should_do_markup(file)
- self._lastlen = 0
- self._chars_on_current_line = 0
- self._width_of_current_line = 0
- @property
- def fullwidth(self):
- if hasattr(self, '_terminal_width'):
- return self._terminal_width
- return get_terminal_width()
- @fullwidth.setter
- def fullwidth(self, value):
- self._terminal_width = value
- @property
- def chars_on_current_line(self):
- """Return the number of characters written so far in the current line.
- Please note that this count does not produce correct results after a reline() call,
- see #164.
- .. versionadded:: 1.5.0
- :rtype: int
- """
- return self._chars_on_current_line
- @property
- def width_of_current_line(self):
- """Return an estimate of the width so far in the current line.
- .. versionadded:: 1.6.0
- :rtype: int
- """
- return self._width_of_current_line
- def _escaped(self, text, esc):
- if esc and self.hasmarkup:
- text = (''.join(['\x1b[%sm' % cod for cod in esc]) +
- text +'\x1b[0m')
- return text
- def markup(self, text, **kw):
- esc = []
- for name in kw:
- if name not in self._esctable:
- raise ValueError("unknown markup: %r" %(name,))
- if kw[name]:
- esc.append(self._esctable[name])
- return self._escaped(text, tuple(esc))
- def sep(self, sepchar, title=None, fullwidth=None, **kw):
- if fullwidth is None:
- fullwidth = self.fullwidth
- # the goal is to have the line be as long as possible
- # under the condition that len(line) <= fullwidth
- if sys.platform == "win32":
- # if we print in the last column on windows we are on a
- # new line but there is no way to verify/neutralize this
- # (we may not know the exact line width)
- # so let's be defensive to avoid empty lines in the output
- fullwidth -= 1
- if title is not None:
- # we want 2 + 2*len(fill) + len(title) <= fullwidth
- # i.e. 2 + 2*len(sepchar)*N + len(title) <= fullwidth
- # 2*len(sepchar)*N <= fullwidth - len(title) - 2
- # N <= (fullwidth - len(title) - 2) // (2*len(sepchar))
- N = max((fullwidth - len(title) - 2) // (2*len(sepchar)), 1)
- fill = sepchar * N
- line = "%s %s %s" % (fill, title, fill)
- else:
- # we want len(sepchar)*N <= fullwidth
- # i.e. N <= fullwidth // len(sepchar)
- line = sepchar * (fullwidth // len(sepchar))
- # in some situations there is room for an extra sepchar at the right,
- # in particular if we consider that with a sepchar like "_ " the
- # trailing space is not important at the end of the line
- if len(line) + len(sepchar.rstrip()) <= fullwidth:
- line += sepchar.rstrip()
- self.line(line, **kw)
- def write(self, msg, **kw):
- if msg:
- if not isinstance(msg, (bytes, text)):
- msg = text(msg)
- self._update_chars_on_current_line(msg)
- if self.hasmarkup and kw:
- markupmsg = self.markup(msg, **kw)
- else:
- markupmsg = msg
- write_out(self._file, markupmsg)
- def _update_chars_on_current_line(self, text_or_bytes):
- newline = b'\n' if isinstance(text_or_bytes, bytes) else '\n'
- current_line = text_or_bytes.rsplit(newline, 1)[-1]
- if isinstance(current_line, bytes):
- current_line = current_line.decode('utf-8', errors='replace')
- if newline in text_or_bytes:
- self._chars_on_current_line = len(current_line)
- self._width_of_current_line = get_line_width(current_line)
- else:
- self._chars_on_current_line += len(current_line)
- self._width_of_current_line += get_line_width(current_line)
- def line(self, s='', **kw):
- self.write(s, **kw)
- self._checkfill(s)
- self.write('\n')
- def reline(self, line, **kw):
- if not self.hasmarkup:
- raise ValueError("cannot use rewrite-line without terminal")
- self.write(line, **kw)
- self._checkfill(line)
- self.write('\r')
- self._lastlen = len(line)
- def _checkfill(self, line):
- diff2last = self._lastlen - len(line)
- if diff2last > 0:
- self.write(" " * diff2last)
- class Win32ConsoleWriter(TerminalWriter):
- def write(self, msg, **kw):
- if msg:
- if not isinstance(msg, (bytes, text)):
- msg = text(msg)
- self._update_chars_on_current_line(msg)
- oldcolors = None
- if self.hasmarkup and kw:
- handle = GetStdHandle(STD_OUTPUT_HANDLE)
- oldcolors = GetConsoleInfo(handle).wAttributes
- default_bg = oldcolors & 0x00F0
- attr = default_bg
- if kw.pop('bold', False):
- attr |= FOREGROUND_INTENSITY
- if kw.pop('red', False):
- attr |= FOREGROUND_RED
- elif kw.pop('blue', False):
- attr |= FOREGROUND_BLUE
- elif kw.pop('green', False):
- attr |= FOREGROUND_GREEN
- elif kw.pop('yellow', False):
- attr |= FOREGROUND_GREEN|FOREGROUND_RED
- else:
- attr |= oldcolors & 0x0007
- SetConsoleTextAttribute(handle, attr)
- write_out(self._file, msg)
- if oldcolors:
- SetConsoleTextAttribute(handle, oldcolors)
- class WriteFile(object):
- def __init__(self, writemethod, encoding=None):
- self.encoding = encoding
- self._writemethod = writemethod
- def write(self, data):
- if self.encoding:
- data = data.encode(self.encoding, "replace")
- self._writemethod(data)
- def flush(self):
- return
- if win32_and_ctypes:
- TerminalWriter = Win32ConsoleWriter
- import ctypes
- from ctypes import wintypes
- # ctypes access to the Windows console
- STD_OUTPUT_HANDLE = -11
- STD_ERROR_HANDLE = -12
- FOREGROUND_BLACK = 0x0000 # black text
- FOREGROUND_BLUE = 0x0001 # text color contains blue.
- FOREGROUND_GREEN = 0x0002 # text color contains green.
- FOREGROUND_RED = 0x0004 # text color contains red.
- FOREGROUND_WHITE = 0x0007
- FOREGROUND_INTENSITY = 0x0008 # text color is intensified.
- BACKGROUND_BLACK = 0x0000 # background color black
- BACKGROUND_BLUE = 0x0010 # background color contains blue.
- BACKGROUND_GREEN = 0x0020 # background color contains green.
- BACKGROUND_RED = 0x0040 # background color contains red.
- BACKGROUND_WHITE = 0x0070
- BACKGROUND_INTENSITY = 0x0080 # background color is intensified.
- SHORT = ctypes.c_short
- class COORD(ctypes.Structure):
- _fields_ = [('X', SHORT),
- ('Y', SHORT)]
- class SMALL_RECT(ctypes.Structure):
- _fields_ = [('Left', SHORT),
- ('Top', SHORT),
- ('Right', SHORT),
- ('Bottom', SHORT)]
- class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):
- _fields_ = [('dwSize', COORD),
- ('dwCursorPosition', COORD),
- ('wAttributes', wintypes.WORD),
- ('srWindow', SMALL_RECT),
- ('dwMaximumWindowSize', COORD)]
- _GetStdHandle = ctypes.windll.kernel32.GetStdHandle
- _GetStdHandle.argtypes = [wintypes.DWORD]
- _GetStdHandle.restype = wintypes.HANDLE
- def GetStdHandle(kind):
- return _GetStdHandle(kind)
- SetConsoleTextAttribute = ctypes.windll.kernel32.SetConsoleTextAttribute
- SetConsoleTextAttribute.argtypes = [wintypes.HANDLE, wintypes.WORD]
- SetConsoleTextAttribute.restype = wintypes.BOOL
- _GetConsoleScreenBufferInfo = \
- ctypes.windll.kernel32.GetConsoleScreenBufferInfo
- _GetConsoleScreenBufferInfo.argtypes = [wintypes.HANDLE,
- ctypes.POINTER(CONSOLE_SCREEN_BUFFER_INFO)]
- _GetConsoleScreenBufferInfo.restype = wintypes.BOOL
- def GetConsoleInfo(handle):
- info = CONSOLE_SCREEN_BUFFER_INFO()
- _GetConsoleScreenBufferInfo(handle, ctypes.byref(info))
- return info
- def _getdimensions():
- handle = GetStdHandle(STD_OUTPUT_HANDLE)
- info = GetConsoleInfo(handle)
- # Substract one from the width, otherwise the cursor wraps
- # and the ending \n causes an empty line to display.
- return info.dwSize.Y, info.dwSize.X - 1
- def write_out(fil, msg):
- # XXX sometimes "msg" is of type bytes, sometimes text which
- # complicates the situation. Should we try to enforce unicode?
- try:
- # on py27 and above writing out to sys.stdout with an encoding
- # should usually work for unicode messages (if the encoding is
- # capable of it)
- fil.write(msg)
- except UnicodeEncodeError:
- # on py26 it might not work because stdout expects bytes
- if fil.encoding:
- try:
- fil.write(msg.encode(fil.encoding))
- except UnicodeEncodeError:
- # it might still fail if the encoding is not capable
- pass
- else:
- fil.flush()
- return
- # fallback: escape all unicode characters
- msg = msg.encode("unicode-escape").decode("ascii")
- fil.write(msg)
- fil.flush()
|