Coverage for /opt/homebrew/lib/python3.11/site-packages/_pytest/_io/terminalwriter.py: 62%

120 statements  

« prev     ^ index     » next       coverage.py v7.2.3, created at 2023-05-04 13:14 +0700

1"""Helper functions for writing to terminals and files.""" 

2import os 

3import shutil 

4import sys 

5from typing import Optional 

6from typing import Sequence 

7from typing import TextIO 

8 

9from .wcwidth import wcswidth 

10from _pytest.compat import final 

11 

12 

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

14 

15 

16def 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 

26def 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 

41class 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 )