aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/python/pytest/py3/_pytest/_io/terminalwriter.py
blob: 379035d858c92b405c02d13a2bd254ce131d2e5e (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
"""Helper functions for writing to terminals and files."""
import os
import shutil
import sys
from typing import Optional
from typing import Sequence
from typing import TextIO

from .wcwidth import wcswidth
from _pytest.compat import final


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


def get_terminal_width() -> int:
    width, _ = shutil.get_terminal_size(fallback=(80, 24))

    # The Windows get_terminal_size may be bogus, let's sanify a bit.
    if width < 40:
        width = 80

    return width


def should_do_markup(file: TextIO) -> bool:
    if os.environ.get("PY_COLORS") == "1":
        return True
    if os.environ.get("PY_COLORS") == "0":
        return False
    if "NO_COLOR" in os.environ:
        return False
    if "FORCE_COLOR" in os.environ:
        return True
    return (
        hasattr(file, "isatty") and file.isatty() and os.environ.get("TERM") != "dumb"
    )


@final
class TerminalWriter:
    _esctable = dict(
        black=30,
        red=31,
        green=32,
        yellow=33,
        blue=34,
        purple=35,
        cyan=36,
        white=37,
        Black=40,
        Red=41,
        Green=42,
        Yellow=43,
        Blue=44,
        Purple=45,
        Cyan=46,
        White=47,
        bold=1,
        light=2,
        blink=5,
        invert=7,
    )

    def __init__(self, file: Optional[TextIO] = None) -> None:
        if file is None:
            file = sys.stdout
        if hasattr(file, "isatty") and file.isatty() and sys.platform == "win32":
            try:
                import colorama
            except ImportError:
                pass
            else:
                file = colorama.AnsiToWin32(file).stream
                assert file is not None
        self._file = file
        self.hasmarkup = should_do_markup(file)
        self._current_line = ""
        self._terminal_width: Optional[int] = None
        self.code_highlight = True

    @property
    def fullwidth(self) -> int:
        if self._terminal_width is not None:
            return self._terminal_width
        return get_terminal_width()

    @fullwidth.setter
    def fullwidth(self, value: int) -> None:
        self._terminal_width = value

    @property
    def width_of_current_line(self) -> int:
        """Return an estimate of the width so far in the current line."""
        return wcswidth(self._current_line)

    def markup(self, text: str, **markup: bool) -> str:
        for name in markup:
            if name not in self._esctable:
                raise ValueError(f"unknown markup: {name!r}")
        if self.hasmarkup:
            esc = [self._esctable[name] for name, on in markup.items() if on]
            if esc:
                text = "".join("\x1b[%sm" % cod for cod in esc) + text + "\x1b[0m"
        return text

    def sep(
        self,
        sepchar: str,
        title: Optional[str] = None,
        fullwidth: Optional[int] = None,
        **markup: bool,
    ) -> None:
        if fullwidth is None:
            fullwidth = self.fullwidth
        # The goal is to have the line be as long as possible
        # under the condition that len(line) <= fullwidth.
        if sys.platform == "win32":
            # If we print in the last column on windows we are on a
            # new line but there is no way to verify/neutralize this
            # (we may not know the exact line width).
            # So let's be defensive to avoid empty lines in the output.
            fullwidth -= 1
        if title is not None:
            # we want 2 + 2*len(fill) + len(title) <= fullwidth
            # i.e.    2 + 2*len(sepchar)*N + len(title) <= fullwidth
            #         2*len(sepchar)*N <= fullwidth - len(title) - 2
            #         N <= (fullwidth - len(title) - 2) // (2*len(sepchar))
            N = max((fullwidth - len(title) - 2) // (2 * len(sepchar)), 1)
            fill = sepchar * N
            line = f"{fill} {title} {fill}"
        else:
            # we want len(sepchar)*N <= fullwidth
            # i.e.    N <= fullwidth // len(sepchar)
            line = sepchar * (fullwidth // len(sepchar))
        # In some situations there is room for an extra sepchar at the right,
        # in particular if we consider that with a sepchar like "_ " the
        # trailing space is not important at the end of the line.
        if len(line) + len(sepchar.rstrip()) <= fullwidth:
            line += sepchar.rstrip()

        self.line(line, **markup)

    def write(self, msg: str, *, flush: bool = False, **markup: bool) -> None:
        if msg:
            current_line = msg.rsplit("\n", 1)[-1]
            if "\n" in msg:
                self._current_line = current_line
            else:
                self._current_line += current_line

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

            try:
                self._file.write(msg)
            except UnicodeEncodeError:
                # Some environments don't support printing general Unicode
                # strings, due to misconfiguration or otherwise; in that case,
                # print the string escaped to ASCII.
                # When the Unicode situation improves we should consider
                # letting the error propagate instead of masking it (see #7475
                # for one brief attempt).
                msg = msg.encode("unicode-escape").decode("ascii")
                self._file.write(msg)

            if flush:
                self.flush()

    def line(self, s: str = "", **markup: bool) -> None:
        self.write(s, **markup)
        self.write("\n")

    def flush(self) -> None:
        self._file.flush()

    def _write_source(self, lines: Sequence[str], indents: Sequence[str] = ()) -> None:
        """Write lines of source code possibly highlighted.

        Keeping this private for now because the API is clunky. We should discuss how
        to evolve the terminal writer so we can have more precise color support, for example
        being able to write part of a line in one color and the rest in another, and so on.
        """
        if indents and len(indents) != len(lines):
            raise ValueError(
                "indents size ({}) should have same size as lines ({})".format(
                    len(indents), len(lines)
                )
            )
        if not indents:
            indents = [""] * len(lines)
        source = "\n".join(lines)
        new_lines = self._highlight(source).splitlines()
        for indent, new_line in zip(indents, new_lines):
            self.line(indent + new_line)

    def _highlight(self, source: str) -> str:
        """Highlight the given source code if we have markup support."""
        from _pytest.config.exceptions import UsageError

        if not self.hasmarkup or not self.code_highlight:
            return source
        try:
            from pygments.formatters.terminal import TerminalFormatter
            from pygments.lexers.python import PythonLexer
            from pygments import highlight
            import pygments.util
        except ImportError:
            return source
        else:
            try:
                highlighted: str = highlight(
                    source,
                    PythonLexer(),
                    TerminalFormatter(
                        bg=os.getenv("PYTEST_THEME_MODE", "dark"),
                        style=os.getenv("PYTEST_THEME"),
                    ),
                )
                return highlighted
            except pygments.util.ClassNotFound:
                raise UsageError(
                    "PYTEST_THEME environment variable had an invalid value: '{}'. "
                    "Only valid pygment styles are allowed.".format(
                        os.getenv("PYTEST_THEME")
                    )
                )
            except pygments.util.OptionError:
                raise UsageError(
                    "PYTEST_THEME_MODE environment variable had an invalid value: '{}'. "
                    "The only allowed values are 'dark' and 'light'.".format(
                        os.getenv("PYTEST_THEME_MODE")
                    )
                )