terminalwriter.py 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233
  1. """Helper functions for writing to terminals and files."""
  2. import os
  3. import shutil
  4. import sys
  5. from typing import Optional
  6. from typing import Sequence
  7. from typing import TextIO
  8. from .wcwidth import wcswidth
  9. from _pytest.compat import final
  10. # This code was initially copied from py 1.8.1, file _io/terminalwriter.py.
  11. def get_terminal_width() -> int:
  12. width, _ = shutil.get_terminal_size(fallback=(80, 24))
  13. # The Windows get_terminal_size may be bogus, let's sanify a bit.
  14. if width < 40:
  15. width = 80
  16. return width
  17. def should_do_markup(file: TextIO) -> bool:
  18. if os.environ.get("PY_COLORS") == "1":
  19. return True
  20. if os.environ.get("PY_COLORS") == "0":
  21. return False
  22. if "NO_COLOR" in os.environ:
  23. return False
  24. if "FORCE_COLOR" in os.environ:
  25. return True
  26. return (
  27. hasattr(file, "isatty") and file.isatty() and os.environ.get("TERM") != "dumb"
  28. )
  29. @final
  30. class TerminalWriter:
  31. _esctable = dict(
  32. black=30,
  33. red=31,
  34. green=32,
  35. yellow=33,
  36. blue=34,
  37. purple=35,
  38. cyan=36,
  39. white=37,
  40. Black=40,
  41. Red=41,
  42. Green=42,
  43. Yellow=43,
  44. Blue=44,
  45. Purple=45,
  46. Cyan=46,
  47. White=47,
  48. bold=1,
  49. light=2,
  50. blink=5,
  51. invert=7,
  52. )
  53. def __init__(self, file: Optional[TextIO] = None) -> None:
  54. if file is None:
  55. file = sys.stdout
  56. if hasattr(file, "isatty") and file.isatty() and sys.platform == "win32":
  57. try:
  58. import colorama
  59. except ImportError:
  60. pass
  61. else:
  62. file = colorama.AnsiToWin32(file).stream
  63. assert file is not None
  64. self._file = file
  65. self.hasmarkup = should_do_markup(file)
  66. self._current_line = ""
  67. self._terminal_width: Optional[int] = None
  68. self.code_highlight = True
  69. @property
  70. def fullwidth(self) -> int:
  71. if self._terminal_width is not None:
  72. return self._terminal_width
  73. return get_terminal_width()
  74. @fullwidth.setter
  75. def fullwidth(self, value: int) -> None:
  76. self._terminal_width = value
  77. @property
  78. def width_of_current_line(self) -> int:
  79. """Return an estimate of the width so far in the current line."""
  80. return wcswidth(self._current_line)
  81. def markup(self, text: str, **markup: bool) -> str:
  82. for name in markup:
  83. if name not in self._esctable:
  84. raise ValueError(f"unknown markup: {name!r}")
  85. if self.hasmarkup:
  86. esc = [self._esctable[name] for name, on in markup.items() if on]
  87. if esc:
  88. text = "".join("\x1b[%sm" % cod for cod in esc) + text + "\x1b[0m"
  89. return text
  90. def sep(
  91. self,
  92. sepchar: str,
  93. title: Optional[str] = None,
  94. fullwidth: Optional[int] = None,
  95. **markup: bool,
  96. ) -> None:
  97. if fullwidth is None:
  98. fullwidth = self.fullwidth
  99. # The goal is to have the line be as long as possible
  100. # under the condition that len(line) <= fullwidth.
  101. if sys.platform == "win32":
  102. # If we print in the last column on windows we are on a
  103. # new line but there is no way to verify/neutralize this
  104. # (we may not know the exact line width).
  105. # So let's be defensive to avoid empty lines in the output.
  106. fullwidth -= 1
  107. if title is not None:
  108. # we want 2 + 2*len(fill) + len(title) <= fullwidth
  109. # i.e. 2 + 2*len(sepchar)*N + len(title) <= fullwidth
  110. # 2*len(sepchar)*N <= fullwidth - len(title) - 2
  111. # N <= (fullwidth - len(title) - 2) // (2*len(sepchar))
  112. N = max((fullwidth - len(title) - 2) // (2 * len(sepchar)), 1)
  113. fill = sepchar * N
  114. line = f"{fill} {title} {fill}"
  115. else:
  116. # we want len(sepchar)*N <= fullwidth
  117. # i.e. N <= fullwidth // len(sepchar)
  118. line = sepchar * (fullwidth // len(sepchar))
  119. # In some situations there is room for an extra sepchar at the right,
  120. # in particular if we consider that with a sepchar like "_ " the
  121. # trailing space is not important at the end of the line.
  122. if len(line) + len(sepchar.rstrip()) <= fullwidth:
  123. line += sepchar.rstrip()
  124. self.line(line, **markup)
  125. def write(self, msg: str, *, flush: bool = False, **markup: bool) -> None:
  126. if msg:
  127. current_line = msg.rsplit("\n", 1)[-1]
  128. if "\n" in msg:
  129. self._current_line = current_line
  130. else:
  131. self._current_line += current_line
  132. msg = self.markup(msg, **markup)
  133. try:
  134. self._file.write(msg)
  135. except UnicodeEncodeError:
  136. # Some environments don't support printing general Unicode
  137. # strings, due to misconfiguration or otherwise; in that case,
  138. # print the string escaped to ASCII.
  139. # When the Unicode situation improves we should consider
  140. # letting the error propagate instead of masking it (see #7475
  141. # for one brief attempt).
  142. msg = msg.encode("unicode-escape").decode("ascii")
  143. self._file.write(msg)
  144. if flush:
  145. self.flush()
  146. def line(self, s: str = "", **markup: bool) -> None:
  147. self.write(s, **markup)
  148. self.write("\n")
  149. def flush(self) -> None:
  150. self._file.flush()
  151. def _write_source(self, lines: Sequence[str], indents: Sequence[str] = ()) -> None:
  152. """Write lines of source code possibly highlighted.
  153. Keeping this private for now because the API is clunky. We should discuss how
  154. to evolve the terminal writer so we can have more precise color support, for example
  155. being able to write part of a line in one color and the rest in another, and so on.
  156. """
  157. if indents and len(indents) != len(lines):
  158. raise ValueError(
  159. "indents size ({}) should have same size as lines ({})".format(
  160. len(indents), len(lines)
  161. )
  162. )
  163. if not indents:
  164. indents = [""] * len(lines)
  165. source = "\n".join(lines)
  166. new_lines = self._highlight(source).splitlines()
  167. for indent, new_line in zip(indents, new_lines):
  168. self.line(indent + new_line)
  169. def _highlight(self, source: str) -> str:
  170. """Highlight the given source code if we have markup support."""
  171. from _pytest.config.exceptions import UsageError
  172. if not self.hasmarkup or not self.code_highlight:
  173. return source
  174. try:
  175. from pygments.formatters.terminal import TerminalFormatter
  176. from pygments.lexers.python import PythonLexer
  177. from pygments import highlight
  178. import pygments.util
  179. except ImportError:
  180. return source
  181. else:
  182. try:
  183. highlighted: str = highlight(
  184. source,
  185. PythonLexer(),
  186. TerminalFormatter(
  187. bg=os.getenv("PYTEST_THEME_MODE", "dark"),
  188. style=os.getenv("PYTEST_THEME"),
  189. ),
  190. )
  191. return highlighted
  192. except pygments.util.ClassNotFound:
  193. raise UsageError(
  194. "PYTEST_THEME environment variable had an invalid value: '{}'. "
  195. "Only valid pygment styles are allowed.".format(
  196. os.getenv("PYTEST_THEME")
  197. )
  198. )
  199. except pygments.util.OptionError:
  200. raise UsageError(
  201. "PYTEST_THEME_MODE environment variable had an invalid value: '{}'. "
  202. "The only allowed values are 'dark' and 'light'.".format(
  203. os.getenv("PYTEST_THEME_MODE")
  204. )
  205. )