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

9 from .wcwidth import wcswidth

10 from _pytest.compat import final

11

12

13# This code was initially copied from py 1.8.1, file _io/terminalwriter.py.

14

15

16 def get_terminal_width() -> int:

17 width, _ = shutil.get_terminal_size(fallback=(80, 24))

18

19 # The Windows get_terminal_size may be bogus, let's sanify a bit.

20 if width < 40:

21 width = 80

22

23 return width

24

25

26 def should_do_markup(file: TextIO) -> bool:

27 if os.environ.get("PY_COLORS") == "1":

28 return True

29 if os.environ.get("PY_COLORS") == "0":

30 return False

31 if "NO_COLOR" in os.environ:

32 return False

33 if "FORCE_COLOR" in os.environ:

34 return True

35 return (

36 hasattr(file, "isatty") and file.isatty() and os.environ.get("TERM") != "dumb"

37 )

38

39

40@final

41 class TerminalWriter:

42 _esctable = dict(

43 black=30,

44 red=31,

45 green=32,

46 yellow=33,

47 blue=34,

48 purple=35,

49 cyan=36,

50 white=37,

51 Black=40,

52 Red=41,

53 Green=42,

54 Yellow=43,

55 Blue=44,

56 Purple=45,

57 Cyan=46,

58 White=47,

59 bold=1,

60 light=2,

61 blink=5,

62 invert=7,

63 )

64

65 def __init__(self, file: Optional[TextIO] = None) -> None:

66 if file is None:

67 file = sys.stdout

68 if hasattr(file, "isatty") and file.isatty() and sys.platform == "win32":

69 try:

70 import colorama

71 except ImportError:

72 pass

73 else:

74 file = colorama.AnsiToWin32(file).stream

75 assert file is not None

76 self._file = file

77 self.hasmarkup = should_do_markup(file)

78 self._current_line = ""

79 self._terminal_width: Optional[int] = None

80 self.code_highlight = True

81

82 @property

83 def fullwidth(self) -> int:

84 if self._terminal_width is not None:

85 return self._terminal_width

86 return get_terminal_width()

87

88 @fullwidth.setter

89 def fullwidth(self, value: int) -> None:

90 self._terminal_width = value

91

92 @property

93 def width_of_current_line(self) -> int:

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

95 return wcswidth(self._current_line)

96

97 def markup(self, text: str, **markup: bool) -> str:

98 for name in markup:

99 if name not in self._esctable:

100 raise ValueError(f"unknown markup: {name!r}")

101 if self.hasmarkup:

102 esc = [self._esctable[name] for name, on in markup.items() if on]

103 if esc:

104 text = "".join("\x1b[%sm" % cod for cod in esc) + text + "\x1b[0m"

105 return text

106

107 def sep(

108 self,

109 sepchar: str,

110 title: Optional[str] = None,

111 fullwidth: Optional[int] = None,

112 **markup: bool,

113 ) -> None:

114 if fullwidth is None:

115 fullwidth = self.fullwidth

116 # The goal is to have the line be as long as possible

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

118 if sys.platform == "win32":

119 # If we print in the last column on windows we are on a

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

121 # (we may not know the exact line width).

122 # So let's be defensive to avoid empty lines in the output.

123 fullwidth -= 1

124 if title is not None:

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

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

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

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

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

130 fill = sepchar * N

131 line = f"{fill} {title} {fill}"

132 else:

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

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

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

136 # In some situations there is room for an extra sepchar at the right,

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

138 # trailing space is not important at the end of the line.

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

140 line += sepchar.rstrip()

141

142 self.line(line, **markup)

143

144 def write(self, msg: str, *, flush: bool = False, **markup: bool) -> None:

145 if msg:

146 current_line = msg.rsplit("\n", 1)[-1]

147 if "\n" in msg:

148 self._current_line = current_line

149 else:

150 self._current_line += current_line

151

152 msg = self.markup(msg, **markup)

153

154 try:

155 self._file.write(msg)

156 except UnicodeEncodeError:

157 # Some environments don't support printing general Unicode

158 # strings, due to misconfiguration or otherwise; in that case,

159 # print the string escaped to ASCII.

160 # When the Unicode situation improves we should consider

161 # letting the error propagate instead of masking it (see #7475

162 # for one brief attempt).

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

164 self._file.write(msg)

165

166 if flush:

167 self.flush()

168

169 def line(self, s: str = "", **markup: bool) -> None:

170 self.write(s, **markup)

171 self.write("\n")

172

173 def flush(self) -> None:

174 self._file.flush()

175

176 def _write_source(self, lines: Sequence[str], indents: Sequence[str] = ()) -> None:

177 """Write lines of source code possibly highlighted.

178

179 Keeping this private for now because the API is clunky. We should discuss how

180 to evolve the terminal writer so we can have more precise color support, for example

181 being able to write part of a line in one color and the rest in another, and so on.

182 """

183 if indents and len(indents) != len(lines):

184 raise ValueError(

185 "indents size ({}) should have same size as lines ({})".format(

186 len(indents), len(lines)

187 )

188 )

189 if not indents:

190 indents = [""] * len(lines)

191 source = "\n".join(lines)

192 new_lines = self._highlight(source).splitlines()

193 for indent, new_line in zip(indents, new_lines):

194 self.line(indent + new_line)

195

196 def _highlight(self, source: str) -> str:

197 """Highlight the given source code if we have markup support."""

198 from _pytest.config.exceptions import UsageError

199

200 if not self.hasmarkup or not self.code_highlight:

201 return source

202 try:

203 from pygments.formatters.terminal import TerminalFormatter

204 from pygments.lexers.python import PythonLexer

205 from pygments import highlight

206 import pygments.util

207 except ImportError:

208 return source

209 else:

210 try:

211 highlighted: str = highlight(

212 source,

213 PythonLexer(),

214 TerminalFormatter(

215 bg=os.getenv("PYTEST_THEME_MODE", "dark"),

216 style=os.getenv("PYTEST_THEME"),

217 ),

218 )

219 return highlighted

220 except pygments.util.ClassNotFound:

221 raise UsageError(

222 "PYTEST_THEME environment variable had an invalid value: '{}'. "

223 "Only valid pygment styles are allowed.".format(

224 os.getenv("PYTEST_THEME")

225 )

226 )

227 except pygments.util.OptionError:

228 raise UsageError(

229 "PYTEST_THEME_MODE environment variable had an invalid value: '{}'. "

230 "The only allowed values are 'dark' and 'light'.".format(

231 os.getenv("PYTEST_THEME_MODE")

232 )

233 )