Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2 

3Helper functions for writing to terminals and files. 

4 

5""" 

6 

7 

8import sys, os, unicodedata 

9import py 

10py3k = sys.version_info[0] >= 3 

11py33 = sys.version_info >= (3, 3) 

12from py.builtin import text, bytes 

13 

14win32_and_ctypes = False 

15colorama = None 

16if sys.platform == "win32": 

17 try: 

18 import colorama 

19 except ImportError: 

20 try: 

21 import ctypes 

22 win32_and_ctypes = True 

23 except ImportError: 

24 pass 

25 

26 

27def _getdimensions(): 

28 if py33: 

29 import shutil 

30 size = shutil.get_terminal_size() 

31 return size.lines, size.columns 

32 else: 

33 import termios, fcntl, struct 

34 call = fcntl.ioctl(1, termios.TIOCGWINSZ, "\000" * 8) 

35 height, width = struct.unpack("hhhh", call)[:2] 

36 return height, width 

37 

38 

39def get_terminal_width(): 

40 width = 0 

41 try: 

42 _, width = _getdimensions() 

43 except py.builtin._sysex: 

44 raise 

45 except: 

46 # pass to fallback below 

47 pass 

48 

49 if width == 0: 

50 # FALLBACK: 

51 # * some exception happened 

52 # * or this is emacs terminal which reports (0,0) 

53 width = int(os.environ.get('COLUMNS', 80)) 

54 

55 # XXX the windows getdimensions may be bogus, let's sanify a bit 

56 if width < 40: 

57 width = 80 

58 return width 

59 

60terminal_width = get_terminal_width() 

61 

62char_width = { 

63 'A': 1, # "Ambiguous" 

64 'F': 2, # Fullwidth 

65 'H': 1, # Halfwidth 

66 'N': 1, # Neutral 

67 'Na': 1, # Narrow 

68 'W': 2, # Wide 

69} 

70 

71 

72def get_line_width(text): 

73 text = unicodedata.normalize('NFC', text) 

74 return sum(char_width.get(unicodedata.east_asian_width(c), 1) for c in text) 

75 

76 

77# XXX unify with _escaped func below 

78def ansi_print(text, esc, file=None, newline=True, flush=False): 

79 if file is None: 

80 file = sys.stderr 

81 text = text.rstrip() 

82 if esc and not isinstance(esc, tuple): 

83 esc = (esc,) 

84 if esc and sys.platform != "win32" and file.isatty(): 

85 text = (''.join(['\x1b[%sm' % cod for cod in esc]) + 

86 text + 

87 '\x1b[0m') # ANSI color code "reset" 

88 if newline: 

89 text += '\n' 

90 

91 if esc and win32_and_ctypes and file.isatty(): 

92 if 1 in esc: 

93 bold = True 

94 esc = tuple([x for x in esc if x != 1]) 

95 else: 

96 bold = False 

97 esctable = {() : FOREGROUND_WHITE, # normal 

98 (31,): FOREGROUND_RED, # red 

99 (32,): FOREGROUND_GREEN, # green 

100 (33,): FOREGROUND_GREEN|FOREGROUND_RED, # yellow 

101 (34,): FOREGROUND_BLUE, # blue 

102 (35,): FOREGROUND_BLUE|FOREGROUND_RED, # purple 

103 (36,): FOREGROUND_BLUE|FOREGROUND_GREEN, # cyan 

104 (37,): FOREGROUND_WHITE, # white 

105 (39,): FOREGROUND_WHITE, # reset 

106 } 

107 attr = esctable.get(esc, FOREGROUND_WHITE) 

108 if bold: 

109 attr |= FOREGROUND_INTENSITY 

110 STD_OUTPUT_HANDLE = -11 

111 STD_ERROR_HANDLE = -12 

112 if file is sys.stderr: 

113 handle = GetStdHandle(STD_ERROR_HANDLE) 

114 else: 

115 handle = GetStdHandle(STD_OUTPUT_HANDLE) 

116 oldcolors = GetConsoleInfo(handle).wAttributes 

117 attr |= (oldcolors & 0x0f0) 

118 SetConsoleTextAttribute(handle, attr) 

119 while len(text) > 32768: 

120 file.write(text[:32768]) 

121 text = text[32768:] 

122 if text: 

123 file.write(text) 

124 SetConsoleTextAttribute(handle, oldcolors) 

125 else: 

126 file.write(text) 

127 

128 if flush: 

129 file.flush() 

130 

131def should_do_markup(file): 

132 if os.environ.get('PY_COLORS') == '1': 

133 return True 

134 if os.environ.get('PY_COLORS') == '0': 

135 return False 

136 return hasattr(file, 'isatty') and file.isatty() \ 

137 and os.environ.get('TERM') != 'dumb' \ 

138 and not (sys.platform.startswith('java') and os._name == 'nt') 

139 

140class TerminalWriter(object): 

141 _esctable = dict(black=30, red=31, green=32, yellow=33, 

142 blue=34, purple=35, cyan=36, white=37, 

143 Black=40, Red=41, Green=42, Yellow=43, 

144 Blue=44, Purple=45, Cyan=46, White=47, 

145 bold=1, light=2, blink=5, invert=7) 

146 

147 # XXX deprecate stringio argument 

148 def __init__(self, file=None, stringio=False, encoding=None): 

149 if file is None: 

150 if stringio: 

151 self.stringio = file = py.io.TextIO() 

152 else: 

153 from sys import stdout as file 

154 elif py.builtin.callable(file) and not ( 

155 hasattr(file, "write") and hasattr(file, "flush")): 

156 file = WriteFile(file, encoding=encoding) 

157 if hasattr(file, "isatty") and file.isatty() and colorama: 

158 file = colorama.AnsiToWin32(file).stream 

159 self.encoding = encoding or getattr(file, 'encoding', "utf-8") 

160 self._file = file 

161 self.hasmarkup = should_do_markup(file) 

162 self._lastlen = 0 

163 self._chars_on_current_line = 0 

164 self._width_of_current_line = 0 

165 

166 @property 

167 def fullwidth(self): 

168 if hasattr(self, '_terminal_width'): 

169 return self._terminal_width 

170 return get_terminal_width() 

171 

172 @fullwidth.setter 

173 def fullwidth(self, value): 

174 self._terminal_width = value 

175 

176 @property 

177 def chars_on_current_line(self): 

178 """Return the number of characters written so far in the current line. 

179 

180 Please note that this count does not produce correct results after a reline() call, 

181 see #164. 

182 

183 .. versionadded:: 1.5.0 

184 

185 :rtype: int 

186 """ 

187 return self._chars_on_current_line 

188 

189 @property 

190 def width_of_current_line(self): 

191 """Return an estimate of the width so far in the current line. 

192 

193 .. versionadded:: 1.6.0 

194 

195 :rtype: int 

196 """ 

197 return self._width_of_current_line 

198 

199 def _escaped(self, text, esc): 

200 if esc and self.hasmarkup: 

201 text = (''.join(['\x1b[%sm' % cod for cod in esc]) + 

202 text +'\x1b[0m') 

203 return text 

204 

205 def markup(self, text, **kw): 

206 esc = [] 

207 for name in kw: 

208 if name not in self._esctable: 

209 raise ValueError("unknown markup: %r" %(name,)) 

210 if kw[name]: 

211 esc.append(self._esctable[name]) 

212 return self._escaped(text, tuple(esc)) 

213 

214 def sep(self, sepchar, title=None, fullwidth=None, **kw): 

215 if fullwidth is None: 

216 fullwidth = self.fullwidth 

217 # the goal is to have the line be as long as possible 

218 # under the condition that len(line) <= fullwidth 

219 if sys.platform == "win32": 

220 # if we print in the last column on windows we are on a 

221 # new line but there is no way to verify/neutralize this 

222 # (we may not know the exact line width) 

223 # so let's be defensive to avoid empty lines in the output 

224 fullwidth -= 1 

225 if title is not None: 

226 # we want 2 + 2*len(fill) + len(title) <= fullwidth 

227 # i.e. 2 + 2*len(sepchar)*N + len(title) <= fullwidth 

228 # 2*len(sepchar)*N <= fullwidth - len(title) - 2 

229 # N <= (fullwidth - len(title) - 2) // (2*len(sepchar)) 

230 N = max((fullwidth - len(title) - 2) // (2*len(sepchar)), 1) 

231 fill = sepchar * N 

232 line = "%s %s %s" % (fill, title, fill) 

233 else: 

234 # we want len(sepchar)*N <= fullwidth 

235 # i.e. N <= fullwidth // len(sepchar) 

236 line = sepchar * (fullwidth // len(sepchar)) 

237 # in some situations there is room for an extra sepchar at the right, 

238 # in particular if we consider that with a sepchar like "_ " the 

239 # trailing space is not important at the end of the line 

240 if len(line) + len(sepchar.rstrip()) <= fullwidth: 

241 line += sepchar.rstrip() 

242 

243 self.line(line, **kw) 

244 

245 def write(self, msg, **kw): 

246 if msg: 

247 if not isinstance(msg, (bytes, text)): 

248 msg = text(msg) 

249 

250 self._update_chars_on_current_line(msg) 

251 

252 if self.hasmarkup and kw: 

253 markupmsg = self.markup(msg, **kw) 

254 else: 

255 markupmsg = msg 

256 write_out(self._file, markupmsg) 

257 

258 def _update_chars_on_current_line(self, text_or_bytes): 

259 newline = b'\n' if isinstance(text_or_bytes, bytes) else '\n' 

260 current_line = text_or_bytes.rsplit(newline, 1)[-1] 

261 if isinstance(current_line, bytes): 

262 current_line = current_line.decode('utf-8', errors='replace') 

263 if newline in text_or_bytes: 

264 self._chars_on_current_line = len(current_line) 

265 self._width_of_current_line = get_line_width(current_line) 

266 else: 

267 self._chars_on_current_line += len(current_line) 

268 self._width_of_current_line += get_line_width(current_line) 

269 

270 def line(self, s='', **kw): 

271 self.write(s, **kw) 

272 self._checkfill(s) 

273 self.write('\n') 

274 

275 def reline(self, line, **kw): 

276 if not self.hasmarkup: 

277 raise ValueError("cannot use rewrite-line without terminal") 

278 self.write(line, **kw) 

279 self._checkfill(line) 

280 self.write('\r') 

281 self._lastlen = len(line) 

282 

283 def _checkfill(self, line): 

284 diff2last = self._lastlen - len(line) 

285 if diff2last > 0: 

286 self.write(" " * diff2last) 

287 

288class Win32ConsoleWriter(TerminalWriter): 

289 def write(self, msg, **kw): 

290 if msg: 

291 if not isinstance(msg, (bytes, text)): 

292 msg = text(msg) 

293 

294 self._update_chars_on_current_line(msg) 

295 

296 oldcolors = None 

297 if self.hasmarkup and kw: 

298 handle = GetStdHandle(STD_OUTPUT_HANDLE) 

299 oldcolors = GetConsoleInfo(handle).wAttributes 

300 default_bg = oldcolors & 0x00F0 

301 attr = default_bg 

302 if kw.pop('bold', False): 

303 attr |= FOREGROUND_INTENSITY 

304 

305 if kw.pop('red', False): 

306 attr |= FOREGROUND_RED 

307 elif kw.pop('blue', False): 

308 attr |= FOREGROUND_BLUE 

309 elif kw.pop('green', False): 

310 attr |= FOREGROUND_GREEN 

311 elif kw.pop('yellow', False): 

312 attr |= FOREGROUND_GREEN|FOREGROUND_RED 

313 else: 

314 attr |= oldcolors & 0x0007 

315 

316 SetConsoleTextAttribute(handle, attr) 

317 write_out(self._file, msg) 

318 if oldcolors: 

319 SetConsoleTextAttribute(handle, oldcolors) 

320 

321class WriteFile(object): 

322 def __init__(self, writemethod, encoding=None): 

323 self.encoding = encoding 

324 self._writemethod = writemethod 

325 

326 def write(self, data): 

327 if self.encoding: 

328 data = data.encode(self.encoding, "replace") 

329 self._writemethod(data) 

330 

331 def flush(self): 

332 return 

333 

334 

335if win32_and_ctypes: 

336 TerminalWriter = Win32ConsoleWriter 

337 import ctypes 

338 from ctypes import wintypes 

339 

340 # ctypes access to the Windows console 

341 STD_OUTPUT_HANDLE = -11 

342 STD_ERROR_HANDLE = -12 

343 FOREGROUND_BLACK = 0x0000 # black text 

344 FOREGROUND_BLUE = 0x0001 # text color contains blue. 

345 FOREGROUND_GREEN = 0x0002 # text color contains green. 

346 FOREGROUND_RED = 0x0004 # text color contains red. 

347 FOREGROUND_WHITE = 0x0007 

348 FOREGROUND_INTENSITY = 0x0008 # text color is intensified. 

349 BACKGROUND_BLACK = 0x0000 # background color black 

350 BACKGROUND_BLUE = 0x0010 # background color contains blue. 

351 BACKGROUND_GREEN = 0x0020 # background color contains green. 

352 BACKGROUND_RED = 0x0040 # background color contains red. 

353 BACKGROUND_WHITE = 0x0070 

354 BACKGROUND_INTENSITY = 0x0080 # background color is intensified. 

355 

356 SHORT = ctypes.c_short 

357 class COORD(ctypes.Structure): 

358 _fields_ = [('X', SHORT), 

359 ('Y', SHORT)] 

360 class SMALL_RECT(ctypes.Structure): 

361 _fields_ = [('Left', SHORT), 

362 ('Top', SHORT), 

363 ('Right', SHORT), 

364 ('Bottom', SHORT)] 

365 class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure): 

366 _fields_ = [('dwSize', COORD), 

367 ('dwCursorPosition', COORD), 

368 ('wAttributes', wintypes.WORD), 

369 ('srWindow', SMALL_RECT), 

370 ('dwMaximumWindowSize', COORD)] 

371 

372 _GetStdHandle = ctypes.windll.kernel32.GetStdHandle 

373 _GetStdHandle.argtypes = [wintypes.DWORD] 

374 _GetStdHandle.restype = wintypes.HANDLE 

375 def GetStdHandle(kind): 

376 return _GetStdHandle(kind) 

377 

378 SetConsoleTextAttribute = ctypes.windll.kernel32.SetConsoleTextAttribute 

379 SetConsoleTextAttribute.argtypes = [wintypes.HANDLE, wintypes.WORD] 

380 SetConsoleTextAttribute.restype = wintypes.BOOL 

381 

382 _GetConsoleScreenBufferInfo = \ 

383 ctypes.windll.kernel32.GetConsoleScreenBufferInfo 

384 _GetConsoleScreenBufferInfo.argtypes = [wintypes.HANDLE, 

385 ctypes.POINTER(CONSOLE_SCREEN_BUFFER_INFO)] 

386 _GetConsoleScreenBufferInfo.restype = wintypes.BOOL 

387 def GetConsoleInfo(handle): 

388 info = CONSOLE_SCREEN_BUFFER_INFO() 

389 _GetConsoleScreenBufferInfo(handle, ctypes.byref(info)) 

390 return info 

391 

392 def _getdimensions(): 

393 handle = GetStdHandle(STD_OUTPUT_HANDLE) 

394 info = GetConsoleInfo(handle) 

395 # Substract one from the width, otherwise the cursor wraps 

396 # and the ending \n causes an empty line to display. 

397 return info.dwSize.Y, info.dwSize.X - 1 

398 

399def write_out(fil, msg): 

400 # XXX sometimes "msg" is of type bytes, sometimes text which 

401 # complicates the situation. Should we try to enforce unicode? 

402 try: 

403 # on py27 and above writing out to sys.stdout with an encoding 

404 # should usually work for unicode messages (if the encoding is 

405 # capable of it) 

406 fil.write(msg) 

407 except UnicodeEncodeError: 

408 # on py26 it might not work because stdout expects bytes 

409 if fil.encoding: 

410 try: 

411 fil.write(msg.encode(fil.encoding)) 

412 except UnicodeEncodeError: 

413 # it might still fail if the encoding is not capable 

414 pass 

415 else: 

416 fil.flush() 

417 return 

418 # fallback: escape all unicode characters 

419 msg = msg.encode("unicode-escape").decode("ascii") 

420 fil.write(msg) 

421 fil.flush()