1# SPDX-License-Identifier: GPL-2.0+
2# Copyright (c) 2011 The Chromium OS Authors.
3#
4
5"""Terminal utilities
6
7This module handles terminal interaction including ANSI color codes.
8"""
9
10from contextlib import contextmanager
11from io import StringIO
12import os
13import re
14import shutil
15import subprocess
16import sys
17
18# Selection of when we want our output to be colored
19COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)
20
21# Initially, we are set up to print to the terminal
22print_test_mode = False
23print_test_list = []
24
25# The length of the last line printed without a newline
26last_print_len = None
27
28# credit:
29# stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python
30ansi_escape = re.compile(r'\x1b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
31
32# True if we are capturing console output
33CAPTURING = False
34
35# Set this to False to disable output-capturing globally
36USE_CAPTURE = True
37
38
39class PrintLine:
40    """A line of text output
41
42    Members:
43        text: Text line that was printed
44        newline: True to output a newline after the text
45        colour: Text colour to use
46    """
47    def __init__(self, text, colour, newline=True, bright=True):
48        self.text = text
49        self.newline = newline
50        self.colour = colour
51        self.bright = bright
52
53    def __eq__(self, other):
54        return (self.text == other.text and
55                self.newline == other.newline and
56                self.colour == other.colour and
57                self.bright == other.bright)
58
59    def __str__(self):
60        return ("newline=%s, colour=%s, bright=%d, text='%s'" %
61                (self.newline, self.colour, self.bright, self.text))
62
63
64def calc_ascii_len(text):
65    """Calculate the length of a string, ignoring any ANSI sequences
66
67    When displayed on a terminal, ANSI sequences don't take any space, so we
68    need to ignore them when calculating the length of a string.
69
70    Args:
71        text: Text to check
72
73    Returns:
74        Length of text, after skipping ANSI sequences
75
76    >>> col = Color(COLOR_ALWAYS)
77    >>> text = col.build(Color.RED, 'abc')
78    >>> len(text)
79    14
80    >>> calc_ascii_len(text)
81    3
82    >>>
83    >>> text += 'def'
84    >>> calc_ascii_len(text)
85    6
86    >>> text += col.build(Color.RED, 'abc')
87    >>> calc_ascii_len(text)
88    9
89    """
90    result = ansi_escape.sub('', text)
91    return len(result)
92
93def trim_ascii_len(text, size):
94    """Trim a string containing ANSI sequences to the given ASCII length
95
96    The string is trimmed with ANSI sequences being ignored for the length
97    calculation.
98
99    >>> col = Color(COLOR_ALWAYS)
100    >>> text = col.build(Color.RED, 'abc')
101    >>> len(text)
102    14
103    >>> calc_ascii_len(trim_ascii_len(text, 4))
104    3
105    >>> calc_ascii_len(trim_ascii_len(text, 2))
106    2
107    >>> text += 'def'
108    >>> calc_ascii_len(trim_ascii_len(text, 4))
109    4
110    >>> text += col.build(Color.RED, 'ghi')
111    >>> calc_ascii_len(trim_ascii_len(text, 7))
112    7
113    """
114    if calc_ascii_len(text) < size:
115        return text
116    pos = 0
117    out = ''
118    left = size
119
120    # Work through each ANSI sequence in turn
121    for m in ansi_escape.finditer(text):
122        # Find the text before the sequence and add it to our string, making
123        # sure it doesn't overflow
124        before = text[pos:m.start()]
125        toadd = before[:left]
126        out += toadd
127
128        # Figure out how much non-ANSI space we have left
129        left -= len(toadd)
130
131        # Add the ANSI sequence and move to the position immediately after it
132        out += m.group()
133        pos = m.start() + len(m.group())
134
135    # Deal with text after the last ANSI sequence
136    after = text[pos:]
137    toadd = after[:left]
138    out += toadd
139
140    return out
141
142
143def tprint(text='', newline=True, colour=None, limit_to_line=False,
144           bright=True, back=None, col=None):
145    """Handle a line of output to the terminal.
146
147    In test mode this is recorded in a list. Otherwise it is output to the
148    terminal.
149
150    Args:
151        text: Text to print
152        newline: True to add a new line at the end of the text
153        colour: Colour to use for the text
154    """
155    global last_print_len
156
157    if print_test_mode:
158        print_test_list.append(PrintLine(text, colour, newline, bright))
159    else:
160        if colour is not None:
161            if not col:
162                col = Color()
163            text = col.build(colour, text, bright=bright, back=back)
164        if newline:
165            print(text)
166            last_print_len = None
167        else:
168            if limit_to_line:
169                cols = shutil.get_terminal_size().columns
170                text = trim_ascii_len(text, cols)
171            print(text, end='', flush=True)
172            last_print_len = calc_ascii_len(text)
173
174def print_clear():
175    """Clear a previously line that was printed with no newline"""
176    global last_print_len
177
178    if last_print_len:
179        if print_test_mode:
180            print_test_list.append(PrintLine(None, None, None, None))
181        else:
182            print('\r%s\r' % (' '* last_print_len), end='', flush=True)
183            last_print_len = None
184
185def set_print_test_mode(enable=True):
186    """Go into test mode, where all printing is recorded"""
187    global print_test_mode
188
189    print_test_mode = enable
190    get_print_test_lines()
191
192def get_print_test_lines():
193    """Get a list of all lines output through tprint()
194
195    Returns:
196        A list of PrintLine objects
197    """
198    global print_test_list
199
200    ret = print_test_list
201    print_test_list = []
202    return ret
203
204def echo_print_test_lines():
205    """Print out the text lines collected"""
206    for line in print_test_list:
207        if line.colour:
208            col = Color()
209            print(col.build(line.colour, line.text), end='')
210        else:
211            print(line.text, end='')
212        if line.newline:
213            print()
214
215def have_terminal():
216    """Check if we have an interactive terminal or not
217
218    Returns:
219        bool: true if an interactive terminal is attached
220    """
221    return os.isatty(sys.stdout.fileno())
222
223
224class Color():
225    """Conditionally wraps text in ANSI color escape sequences."""
226    BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
227    BOLD = -1
228    BRIGHT_START = '\033[1;%d%sm'
229    NORMAL_START = '\033[22;%d%sm'
230    BOLD_START = '\033[1m'
231    BACK_EXTRA = ';%d'
232    RESET = '\033[0m'
233
234    def __init__(self, colored=COLOR_IF_TERMINAL):
235        """Create a new Color object, optionally disabling color output.
236
237        Args:
238          enabled: True if color output should be enabled. If False then this
239            class will not add color codes at all.
240        """
241        try:
242            self._enabled = (colored == COLOR_ALWAYS or
243                    (colored == COLOR_IF_TERMINAL and
244                     os.isatty(sys.stdout.fileno())))
245        except:
246            self._enabled = False
247
248    def enabled(self):
249        """Check if colour is enabled
250
251        Return: True if enabled, else False
252        """
253        return self._enabled
254
255    def start(self, color, bright=True, back=None):
256        """Returns a start color code.
257
258        Args:
259          color: Color to use, .e.g BLACK, RED, etc.
260
261        Returns:
262          If color is enabled, returns an ANSI sequence to start the given
263          color, otherwise returns empty string
264        """
265        if self._enabled:
266            if color == self.BOLD:
267                return self.BOLD_START
268            base = self.BRIGHT_START if bright else self.NORMAL_START
269            extra = self.BACK_EXTRA % (back + 40) if back else ''
270            return base % (color + 30, extra)
271        return ''
272
273    def stop(self):
274        """Returns a stop color code.
275
276        Returns:
277          If color is enabled, returns an ANSI color reset sequence,
278          otherwise returns empty string
279        """
280        if self._enabled:
281            return self.RESET
282        return ''
283
284    def build(self, color, text, bright=True, back=None):
285        """Returns text with conditionally added color escape sequences.
286
287        Keyword arguments:
288          color: Text color -- one of the color constants defined in this
289                  class.
290          text: The text to color.
291
292        Returns:
293          If self._enabled is False, returns the original text. If it's True,
294          returns text with color escape sequences based on the value of
295          color.
296        """
297        if not self._enabled:
298            return text
299        return self.start(color, bright, back) + text + self.RESET
300
301
302# Use this to suppress stdout/stderr output:
303# with terminal.capture() as (stdout, stderr)
304#   ...do something...
305@contextmanager
306def capture():
307    global CAPTURING
308
309    capture_out, capture_err = StringIO(), StringIO()
310    old_out, old_err = sys.stdout, sys.stderr
311    try:
312        CAPTURING = True
313        sys.stdout, sys.stderr = capture_out, capture_err
314        yield capture_out, capture_err
315    finally:
316        sys.stdout, sys.stderr = old_out, old_err
317        CAPTURING = False
318        if not USE_CAPTURE:
319            sys.stdout.write(capture_out.getvalue())
320            sys.stderr.write(capture_err.getvalue())
321
322
323@contextmanager
324def pager():
325    """Simple pager for outputting lots of text
326
327    Usage:
328        with terminal.pager():
329            print(...)
330    """
331    proc = None
332    old_stdout = None
333    try:
334        less = os.getenv('PAGER')
335        if not CAPTURING and less != 'none' and have_terminal():
336            if not less:
337                less = 'less -R --quit-if-one-screen'
338            proc = subprocess.Popen(less, stdin=subprocess.PIPE, text=True,
339                                    shell=True)
340            old_stdout = sys.stdout
341            sys.stdout = proc.stdin
342        yield
343    finally:
344        if proc:
345            sys.stdout = old_stdout
346            proc.communicate()
347