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
« 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
9from .wcwidth import wcswidth
10from _pytest.compat import final
13# This code was initially copied from py 1.8.1, file _io/terminalwriter.py.
16def get_terminal_width() -> int:
17 width, _ = shutil.get_terminal_size(fallback=(80, 24))
19 # The Windows get_terminal_size may be bogus, let's sanify a bit.
20 if width < 40:
21 width = 80
23 return width
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 )
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 )
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
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()
88 @fullwidth.setter
89 def fullwidth(self, value: int) -> None:
90 self._terminal_width = value
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)
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
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()
142 self.line(line, **markup)
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
152 msg = self.markup(msg, **markup)
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)
166 if flush:
167 self.flush()
169 def line(self, s: str = "", **markup: bool) -> None:
170 self.write(s, **markup)
171 self.write("\n")
173 def flush(self) -> None:
174 self._file.flush()
176 def _write_source(self, lines: Sequence[str], indents: Sequence[str] = ()) -> None:
177 """Write lines of source code possibly highlighted.
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)
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
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 )