1""" 2 3Helper functions for writing to terminals and files. 4 5""" 6 7 8import sys, os 9import py 10py3k = sys.version_info[0] >= 3 11from py.builtin import text, bytes 12 13win32_and_ctypes = False 14colorama = None 15if sys.platform == "win32": 16 try: 17 import colorama 18 except ImportError: 19 try: 20 import ctypes 21 win32_and_ctypes = True 22 except ImportError: 23 pass 24 25 26def _getdimensions(): 27 import termios,fcntl,struct 28 call = fcntl.ioctl(1,termios.TIOCGWINSZ,"\000"*8) 29 height,width = struct.unpack( "hhhh", call ) [:2] 30 return height, width 31 32 33def get_terminal_width(): 34 width = 0 35 try: 36 _, width = _getdimensions() 37 except py.builtin._sysex: 38 raise 39 except: 40 # pass to fallback below 41 pass 42 43 if width == 0: 44 # FALLBACK: 45 # * some exception happened 46 # * or this is emacs terminal which reports (0,0) 47 width = int(os.environ.get('COLUMNS', 80)) 48 49 # XXX the windows getdimensions may be bogus, let's sanify a bit 50 if width < 40: 51 width = 80 52 return width 53 54terminal_width = get_terminal_width() 55 56# XXX unify with _escaped func below 57def ansi_print(text, esc, file=None, newline=True, flush=False): 58 if file is None: 59 file = sys.stderr 60 text = text.rstrip() 61 if esc and not isinstance(esc, tuple): 62 esc = (esc,) 63 if esc and sys.platform != "win32" and file.isatty(): 64 text = (''.join(['\x1b[%sm' % cod for cod in esc]) + 65 text + 66 '\x1b[0m') # ANSI color code "reset" 67 if newline: 68 text += '\n' 69 70 if esc and win32_and_ctypes and file.isatty(): 71 if 1 in esc: 72 bold = True 73 esc = tuple([x for x in esc if x != 1]) 74 else: 75 bold = False 76 esctable = {() : FOREGROUND_WHITE, # normal 77 (31,): FOREGROUND_RED, # red 78 (32,): FOREGROUND_GREEN, # green 79 (33,): FOREGROUND_GREEN|FOREGROUND_RED, # yellow 80 (34,): FOREGROUND_BLUE, # blue 81 (35,): FOREGROUND_BLUE|FOREGROUND_RED, # purple 82 (36,): FOREGROUND_BLUE|FOREGROUND_GREEN, # cyan 83 (37,): FOREGROUND_WHITE, # white 84 (39,): FOREGROUND_WHITE, # reset 85 } 86 attr = esctable.get(esc, FOREGROUND_WHITE) 87 if bold: 88 attr |= FOREGROUND_INTENSITY 89 STD_OUTPUT_HANDLE = -11 90 STD_ERROR_HANDLE = -12 91 if file is sys.stderr: 92 handle = GetStdHandle(STD_ERROR_HANDLE) 93 else: 94 handle = GetStdHandle(STD_OUTPUT_HANDLE) 95 oldcolors = GetConsoleInfo(handle).wAttributes 96 attr |= (oldcolors & 0x0f0) 97 SetConsoleTextAttribute(handle, attr) 98 while len(text) > 32768: 99 file.write(text[:32768]) 100 text = text[32768:] 101 if text: 102 file.write(text) 103 SetConsoleTextAttribute(handle, oldcolors) 104 else: 105 file.write(text) 106 107 if flush: 108 file.flush() 109 110def should_do_markup(file): 111 if os.environ.get('PY_COLORS') == '1': 112 return True 113 if os.environ.get('PY_COLORS') == '0': 114 return False 115 return hasattr(file, 'isatty') and file.isatty() \ 116 and os.environ.get('TERM') != 'dumb' \ 117 and not (sys.platform.startswith('java') and os._name == 'nt') 118 119class TerminalWriter(object): 120 _esctable = dict(black=30, red=31, green=32, yellow=33, 121 blue=34, purple=35, cyan=36, white=37, 122 Black=40, Red=41, Green=42, Yellow=43, 123 Blue=44, Purple=45, Cyan=46, White=47, 124 bold=1, light=2, blink=5, invert=7) 125 126 # XXX deprecate stringio argument 127 def __init__(self, file=None, stringio=False, encoding=None): 128 if file is None: 129 if stringio: 130 self.stringio = file = py.io.TextIO() 131 else: 132 from sys import stdout as file 133 elif py.builtin.callable(file) and not ( 134 hasattr(file, "write") and hasattr(file, "flush")): 135 file = WriteFile(file, encoding=encoding) 136 if hasattr(file, "isatty") and file.isatty() and colorama: 137 file = colorama.AnsiToWin32(file).stream 138 self.encoding = encoding or getattr(file, 'encoding', "utf-8") 139 self._file = file 140 self.hasmarkup = should_do_markup(file) 141 self._lastlen = 0 142 self._chars_on_current_line = 0 143 144 @property 145 def fullwidth(self): 146 if hasattr(self, '_terminal_width'): 147 return self._terminal_width 148 return get_terminal_width() 149 150 @fullwidth.setter 151 def fullwidth(self, value): 152 self._terminal_width = value 153 154 @property 155 def chars_on_current_line(self): 156 """Return the number of characters written so far in the current line. 157 158 Please note that this count does not produce correct results after a reline() call, 159 see #164. 160 161 .. versionadded:: 1.5.0 162 163 :rtype: int 164 """ 165 return self._chars_on_current_line 166 167 def _escaped(self, text, esc): 168 if esc and self.hasmarkup: 169 text = (''.join(['\x1b[%sm' % cod for cod in esc]) + 170 text +'\x1b[0m') 171 return text 172 173 def markup(self, text, **kw): 174 esc = [] 175 for name in kw: 176 if name not in self._esctable: 177 raise ValueError("unknown markup: %r" %(name,)) 178 if kw[name]: 179 esc.append(self._esctable[name]) 180 return self._escaped(text, tuple(esc)) 181 182 def sep(self, sepchar, title=None, fullwidth=None, **kw): 183 if fullwidth is None: 184 fullwidth = self.fullwidth 185 # the goal is to have the line be as long as possible 186 # under the condition that len(line) <= fullwidth 187 if sys.platform == "win32": 188 # if we print in the last column on windows we are on a 189 # new line but there is no way to verify/neutralize this 190 # (we may not know the exact line width) 191 # so let's be defensive to avoid empty lines in the output 192 fullwidth -= 1 193 if title is not None: 194 # we want 2 + 2*len(fill) + len(title) <= fullwidth 195 # i.e. 2 + 2*len(sepchar)*N + len(title) <= fullwidth 196 # 2*len(sepchar)*N <= fullwidth - len(title) - 2 197 # N <= (fullwidth - len(title) - 2) // (2*len(sepchar)) 198 N = (fullwidth - len(title) - 2) // (2*len(sepchar)) 199 fill = sepchar * N 200 line = "%s %s %s" % (fill, title, fill) 201 else: 202 # we want len(sepchar)*N <= fullwidth 203 # i.e. N <= fullwidth // len(sepchar) 204 line = sepchar * (fullwidth // len(sepchar)) 205 # in some situations there is room for an extra sepchar at the right, 206 # in particular if we consider that with a sepchar like "_ " the 207 # trailing space is not important at the end of the line 208 if len(line) + len(sepchar.rstrip()) <= fullwidth: 209 line += sepchar.rstrip() 210 211 self.line(line, **kw) 212 213 def write(self, msg, **kw): 214 if msg: 215 if not isinstance(msg, (bytes, text)): 216 msg = text(msg) 217 218 self._update_chars_on_current_line(msg) 219 220 if self.hasmarkup and kw: 221 markupmsg = self.markup(msg, **kw) 222 else: 223 markupmsg = msg 224 write_out(self._file, markupmsg) 225 226 def _update_chars_on_current_line(self, text): 227 fields = text.rsplit('\n', 1) 228 if '\n' in text: 229 self._chars_on_current_line = len(fields[-1]) 230 else: 231 self._chars_on_current_line += len(fields[-1]) 232 233 def line(self, s='', **kw): 234 self.write(s, **kw) 235 self._checkfill(s) 236 self.write('\n') 237 238 def reline(self, line, **kw): 239 if not self.hasmarkup: 240 raise ValueError("cannot use rewrite-line without terminal") 241 self.write(line, **kw) 242 self._checkfill(line) 243 self.write('\r') 244 self._lastlen = len(line) 245 246 def _checkfill(self, line): 247 diff2last = self._lastlen - len(line) 248 if diff2last > 0: 249 self.write(" " * diff2last) 250 251class Win32ConsoleWriter(TerminalWriter): 252 def write(self, msg, **kw): 253 if msg: 254 if not isinstance(msg, (bytes, text)): 255 msg = text(msg) 256 257 self._update_chars_on_current_line(msg) 258 259 oldcolors = None 260 if self.hasmarkup and kw: 261 handle = GetStdHandle(STD_OUTPUT_HANDLE) 262 oldcolors = GetConsoleInfo(handle).wAttributes 263 default_bg = oldcolors & 0x00F0 264 attr = default_bg 265 if kw.pop('bold', False): 266 attr |= FOREGROUND_INTENSITY 267 268 if kw.pop('red', False): 269 attr |= FOREGROUND_RED 270 elif kw.pop('blue', False): 271 attr |= FOREGROUND_BLUE 272 elif kw.pop('green', False): 273 attr |= FOREGROUND_GREEN 274 elif kw.pop('yellow', False): 275 attr |= FOREGROUND_GREEN|FOREGROUND_RED 276 else: 277 attr |= oldcolors & 0x0007 278 279 SetConsoleTextAttribute(handle, attr) 280 write_out(self._file, msg) 281 if oldcolors: 282 SetConsoleTextAttribute(handle, oldcolors) 283 284class WriteFile(object): 285 def __init__(self, writemethod, encoding=None): 286 self.encoding = encoding 287 self._writemethod = writemethod 288 289 def write(self, data): 290 if self.encoding: 291 data = data.encode(self.encoding, "replace") 292 self._writemethod(data) 293 294 def flush(self): 295 return 296 297 298if win32_and_ctypes: 299 TerminalWriter = Win32ConsoleWriter 300 import ctypes 301 from ctypes import wintypes 302 303 # ctypes access to the Windows console 304 STD_OUTPUT_HANDLE = -11 305 STD_ERROR_HANDLE = -12 306 FOREGROUND_BLACK = 0x0000 # black text 307 FOREGROUND_BLUE = 0x0001 # text color contains blue. 308 FOREGROUND_GREEN = 0x0002 # text color contains green. 309 FOREGROUND_RED = 0x0004 # text color contains red. 310 FOREGROUND_WHITE = 0x0007 311 FOREGROUND_INTENSITY = 0x0008 # text color is intensified. 312 BACKGROUND_BLACK = 0x0000 # background color black 313 BACKGROUND_BLUE = 0x0010 # background color contains blue. 314 BACKGROUND_GREEN = 0x0020 # background color contains green. 315 BACKGROUND_RED = 0x0040 # background color contains red. 316 BACKGROUND_WHITE = 0x0070 317 BACKGROUND_INTENSITY = 0x0080 # background color is intensified. 318 319 SHORT = ctypes.c_short 320 class COORD(ctypes.Structure): 321 _fields_ = [('X', SHORT), 322 ('Y', SHORT)] 323 class SMALL_RECT(ctypes.Structure): 324 _fields_ = [('Left', SHORT), 325 ('Top', SHORT), 326 ('Right', SHORT), 327 ('Bottom', SHORT)] 328 class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure): 329 _fields_ = [('dwSize', COORD), 330 ('dwCursorPosition', COORD), 331 ('wAttributes', wintypes.WORD), 332 ('srWindow', SMALL_RECT), 333 ('dwMaximumWindowSize', COORD)] 334 335 _GetStdHandle = ctypes.windll.kernel32.GetStdHandle 336 _GetStdHandle.argtypes = [wintypes.DWORD] 337 _GetStdHandle.restype = wintypes.HANDLE 338 def GetStdHandle(kind): 339 return _GetStdHandle(kind) 340 341 SetConsoleTextAttribute = ctypes.windll.kernel32.SetConsoleTextAttribute 342 SetConsoleTextAttribute.argtypes = [wintypes.HANDLE, wintypes.WORD] 343 SetConsoleTextAttribute.restype = wintypes.BOOL 344 345 _GetConsoleScreenBufferInfo = \ 346 ctypes.windll.kernel32.GetConsoleScreenBufferInfo 347 _GetConsoleScreenBufferInfo.argtypes = [wintypes.HANDLE, 348 ctypes.POINTER(CONSOLE_SCREEN_BUFFER_INFO)] 349 _GetConsoleScreenBufferInfo.restype = wintypes.BOOL 350 def GetConsoleInfo(handle): 351 info = CONSOLE_SCREEN_BUFFER_INFO() 352 _GetConsoleScreenBufferInfo(handle, ctypes.byref(info)) 353 return info 354 355 def _getdimensions(): 356 handle = GetStdHandle(STD_OUTPUT_HANDLE) 357 info = GetConsoleInfo(handle) 358 # Substract one from the width, otherwise the cursor wraps 359 # and the ending \n causes an empty line to display. 360 return info.dwSize.Y, info.dwSize.X - 1 361 362def write_out(fil, msg): 363 # XXX sometimes "msg" is of type bytes, sometimes text which 364 # complicates the situation. Should we try to enforce unicode? 365 try: 366 # on py27 and above writing out to sys.stdout with an encoding 367 # should usually work for unicode messages (if the encoding is 368 # capable of it) 369 fil.write(msg) 370 except UnicodeEncodeError: 371 # on py26 it might not work because stdout expects bytes 372 if fil.encoding: 373 try: 374 fil.write(msg.encode(fil.encoding)) 375 except UnicodeEncodeError: 376 # it might still fail if the encoding is not capable 377 pass 378 else: 379 fil.flush() 380 return 381 # fallback: escape all unicode characters 382 msg = msg.encode("unicode-escape").decode("ascii") 383 fil.write(msg) 384 fil.flush() 385