terminalwriter.py 14 KB

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