You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

422 lines
14KB

  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. return hasattr(file, 'isatty') and file.isatty() \
  117. and os.environ.get('TERM') != 'dumb' \
  118. and not (sys.platform.startswith('java') and os._name == 'nt')
  119. class 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. # XXX deprecate stringio argument
  126. def __init__(self, file=None, stringio=False, encoding=None):
  127. if file is None:
  128. if stringio:
  129. self.stringio = file = py.io.TextIO()
  130. else:
  131. from sys import stdout as file
  132. elif py.builtin.callable(file) and not (
  133. hasattr(file, "write") and hasattr(file, "flush")):
  134. file = WriteFile(file, encoding=encoding)
  135. if hasattr(file, "isatty") and file.isatty() and colorama:
  136. file = colorama.AnsiToWin32(file).stream
  137. self.encoding = encoding or getattr(file, 'encoding', "utf-8")
  138. self._file = file
  139. self.hasmarkup = should_do_markup(file)
  140. self._lastlen = 0
  141. self._chars_on_current_line = 0
  142. self._width_of_current_line = 0
  143. @property
  144. def fullwidth(self):
  145. if hasattr(self, '_terminal_width'):
  146. return self._terminal_width
  147. return get_terminal_width()
  148. @fullwidth.setter
  149. def fullwidth(self, value):
  150. self._terminal_width = value
  151. @property
  152. def chars_on_current_line(self):
  153. """Return the number of characters written so far in the current line.
  154. Please note that this count does not produce correct results after a reline() call,
  155. see #164.
  156. .. versionadded:: 1.5.0
  157. :rtype: int
  158. """
  159. return self._chars_on_current_line
  160. @property
  161. def width_of_current_line(self):
  162. """Return an estimate of the width so far in the current line.
  163. .. versionadded:: 1.6.0
  164. :rtype: int
  165. """
  166. return self._width_of_current_line
  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. def markup(self, text, **kw):
  173. esc = []
  174. for name in kw:
  175. if name not in self._esctable:
  176. raise ValueError("unknown markup: %r" %(name,))
  177. if kw[name]:
  178. esc.append(self._esctable[name])
  179. return self._escaped(text, tuple(esc))
  180. def sep(self, sepchar, title=None, fullwidth=None, **kw):
  181. if fullwidth is None:
  182. fullwidth = self.fullwidth
  183. # the goal is to have the line be as long as possible
  184. # under the condition that len(line) <= fullwidth
  185. if sys.platform == "win32":
  186. # if we print in the last column on windows we are on a
  187. # new line but there is no way to verify/neutralize this
  188. # (we may not know the exact line width)
  189. # so let's be defensive to avoid empty lines in the output
  190. fullwidth -= 1
  191. if title is not None:
  192. # we want 2 + 2*len(fill) + len(title) <= fullwidth
  193. # i.e. 2 + 2*len(sepchar)*N + len(title) <= fullwidth
  194. # 2*len(sepchar)*N <= fullwidth - len(title) - 2
  195. # N <= (fullwidth - len(title) - 2) // (2*len(sepchar))
  196. N = max((fullwidth - len(title) - 2) // (2*len(sepchar)), 1)
  197. fill = sepchar * N
  198. line = "%s %s %s" % (fill, title, fill)
  199. else:
  200. # we want len(sepchar)*N <= fullwidth
  201. # i.e. N <= fullwidth // len(sepchar)
  202. line = sepchar * (fullwidth // len(sepchar))
  203. # in some situations there is room for an extra sepchar at the right,
  204. # in particular if we consider that with a sepchar like "_ " the
  205. # trailing space is not important at the end of the line
  206. if len(line) + len(sepchar.rstrip()) <= fullwidth:
  207. line += sepchar.rstrip()
  208. self.line(line, **kw)
  209. def write(self, msg, **kw):
  210. if msg:
  211. if not isinstance(msg, (bytes, text)):
  212. msg = text(msg)
  213. self._update_chars_on_current_line(msg)
  214. if self.hasmarkup and kw:
  215. markupmsg = self.markup(msg, **kw)
  216. else:
  217. markupmsg = msg
  218. write_out(self._file, markupmsg)
  219. def _update_chars_on_current_line(self, text_or_bytes):
  220. newline = b'\n' if isinstance(text_or_bytes, bytes) else '\n'
  221. current_line = text_or_bytes.rsplit(newline, 1)[-1]
  222. if isinstance(current_line, bytes):
  223. current_line = current_line.decode('utf-8', errors='replace')
  224. if newline in text_or_bytes:
  225. self._chars_on_current_line = len(current_line)
  226. self._width_of_current_line = get_line_width(current_line)
  227. else:
  228. self._chars_on_current_line += len(current_line)
  229. self._width_of_current_line += get_line_width(current_line)
  230. def line(self, s='', **kw):
  231. self.write(s, **kw)
  232. self._checkfill(s)
  233. self.write('\n')
  234. def reline(self, line, **kw):
  235. if not self.hasmarkup:
  236. raise ValueError("cannot use rewrite-line without terminal")
  237. self.write(line, **kw)
  238. self._checkfill(line)
  239. self.write('\r')
  240. self._lastlen = len(line)
  241. def _checkfill(self, line):
  242. diff2last = self._lastlen - len(line)
  243. if diff2last > 0:
  244. self.write(" " * diff2last)
  245. class Win32ConsoleWriter(TerminalWriter):
  246. def write(self, msg, **kw):
  247. if msg:
  248. if not isinstance(msg, (bytes, text)):
  249. msg = text(msg)
  250. self._update_chars_on_current_line(msg)
  251. oldcolors = None
  252. if self.hasmarkup and kw:
  253. handle = GetStdHandle(STD_OUTPUT_HANDLE)
  254. oldcolors = GetConsoleInfo(handle).wAttributes
  255. default_bg = oldcolors & 0x00F0
  256. attr = default_bg
  257. if kw.pop('bold', False):
  258. attr |= FOREGROUND_INTENSITY
  259. if kw.pop('red', False):
  260. attr |= FOREGROUND_RED
  261. elif kw.pop('blue', False):
  262. attr |= FOREGROUND_BLUE
  263. elif kw.pop('green', False):
  264. attr |= FOREGROUND_GREEN
  265. elif kw.pop('yellow', False):
  266. attr |= FOREGROUND_GREEN|FOREGROUND_RED
  267. else:
  268. attr |= oldcolors & 0x0007
  269. SetConsoleTextAttribute(handle, attr)
  270. write_out(self._file, msg)
  271. if oldcolors:
  272. SetConsoleTextAttribute(handle, oldcolors)
  273. class WriteFile(object):
  274. def __init__(self, writemethod, encoding=None):
  275. self.encoding = encoding
  276. self._writemethod = writemethod
  277. def write(self, data):
  278. if self.encoding:
  279. data = data.encode(self.encoding, "replace")
  280. self._writemethod(data)
  281. def flush(self):
  282. return
  283. if win32_and_ctypes:
  284. TerminalWriter = Win32ConsoleWriter
  285. import ctypes
  286. from ctypes import wintypes
  287. # ctypes access to the Windows console
  288. STD_OUTPUT_HANDLE = -11
  289. STD_ERROR_HANDLE = -12
  290. FOREGROUND_BLACK = 0x0000 # black text
  291. FOREGROUND_BLUE = 0x0001 # text color contains blue.
  292. FOREGROUND_GREEN = 0x0002 # text color contains green.
  293. FOREGROUND_RED = 0x0004 # text color contains red.
  294. FOREGROUND_WHITE = 0x0007
  295. FOREGROUND_INTENSITY = 0x0008 # text color is intensified.
  296. BACKGROUND_BLACK = 0x0000 # background color black
  297. BACKGROUND_BLUE = 0x0010 # background color contains blue.
  298. BACKGROUND_GREEN = 0x0020 # background color contains green.
  299. BACKGROUND_RED = 0x0040 # background color contains red.
  300. BACKGROUND_WHITE = 0x0070
  301. BACKGROUND_INTENSITY = 0x0080 # background color is intensified.
  302. SHORT = ctypes.c_short
  303. class COORD(ctypes.Structure):
  304. _fields_ = [('X', SHORT),
  305. ('Y', SHORT)]
  306. class SMALL_RECT(ctypes.Structure):
  307. _fields_ = [('Left', SHORT),
  308. ('Top', SHORT),
  309. ('Right', SHORT),
  310. ('Bottom', SHORT)]
  311. class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):
  312. _fields_ = [('dwSize', COORD),
  313. ('dwCursorPosition', COORD),
  314. ('wAttributes', wintypes.WORD),
  315. ('srWindow', SMALL_RECT),
  316. ('dwMaximumWindowSize', COORD)]
  317. _GetStdHandle = ctypes.windll.kernel32.GetStdHandle
  318. _GetStdHandle.argtypes = [wintypes.DWORD]
  319. _GetStdHandle.restype = wintypes.HANDLE
  320. def GetStdHandle(kind):
  321. return _GetStdHandle(kind)
  322. SetConsoleTextAttribute = ctypes.windll.kernel32.SetConsoleTextAttribute
  323. SetConsoleTextAttribute.argtypes = [wintypes.HANDLE, wintypes.WORD]
  324. SetConsoleTextAttribute.restype = wintypes.BOOL
  325. _GetConsoleScreenBufferInfo = \
  326. ctypes.windll.kernel32.GetConsoleScreenBufferInfo
  327. _GetConsoleScreenBufferInfo.argtypes = [wintypes.HANDLE,
  328. ctypes.POINTER(CONSOLE_SCREEN_BUFFER_INFO)]
  329. _GetConsoleScreenBufferInfo.restype = wintypes.BOOL
  330. def GetConsoleInfo(handle):
  331. info = CONSOLE_SCREEN_BUFFER_INFO()
  332. _GetConsoleScreenBufferInfo(handle, ctypes.byref(info))
  333. return info
  334. def _getdimensions():
  335. handle = GetStdHandle(STD_OUTPUT_HANDLE)
  336. info = GetConsoleInfo(handle)
  337. # Substract one from the width, otherwise the cursor wraps
  338. # and the ending \n causes an empty line to display.
  339. return info.dwSize.Y, info.dwSize.X - 1
  340. def write_out(fil, msg):
  341. # XXX sometimes "msg" is of type bytes, sometimes text which
  342. # complicates the situation. Should we try to enforce unicode?
  343. try:
  344. # on py27 and above writing out to sys.stdout with an encoding
  345. # should usually work for unicode messages (if the encoding is
  346. # capable of it)
  347. fil.write(msg)
  348. except UnicodeEncodeError:
  349. # on py26 it might not work because stdout expects bytes
  350. if fil.encoding:
  351. try:
  352. fil.write(msg.encode(fil.encoding))
  353. except UnicodeEncodeError:
  354. # it might still fail if the encoding is not capable
  355. pass
  356. else:
  357. fil.flush()
  358. return
  359. # fallback: escape all unicode characters
  360. msg = msg.encode("unicode-escape").decode("ascii")
  361. fil.write(msg)
  362. fil.flush()