terminalwriter.py 13 KB

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