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
10import os
11import re
12import shutil
13import sys
14
15# Selection of when we want our output to be colored
16COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)
17
18# Initially, we are set up to print to the terminal
19print_test_mode = False
20print_test_list = []
21
22# The length of the last line printed without a newline
23last_print_len = None
24
25# credit:
26# stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python
27ansi_escape = re.compile(r'\x1b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
28
29class PrintLine:
30    """A line of text output
31
32    Members:
33        text: Text line that was printed
34        newline: True to output a newline after the text
35        colour: Text colour to use
36    """
37    def __init__(self, text, newline, colour):
38        self.text = text
39        self.newline = newline
40        self.colour = colour
41
42    def __str__(self):
43        return 'newline=%s, colour=%s, text=%s' % (self.newline, self.colour,
44                self.text)
45
46def CalcAsciiLen(text):
47    """Calculate the length of a string, ignoring any ANSI sequences
48
49    When displayed on a terminal, ANSI sequences don't take any space, so we
50    need to ignore them when calculating the length of a string.
51
52    Args:
53        text: Text to check
54
55    Returns:
56        Length of text, after skipping ANSI sequences
57
58    >>> col = Color(COLOR_ALWAYS)
59    >>> text = col.Color(Color.RED, 'abc')
60    >>> len(text)
61    14
62    >>> CalcAsciiLen(text)
63    3
64    >>>
65    >>> text += 'def'
66    >>> CalcAsciiLen(text)
67    6
68    >>> text += col.Color(Color.RED, 'abc')
69    >>> CalcAsciiLen(text)
70    9
71    """
72    result = ansi_escape.sub('', text)
73    return len(result)
74
75def TrimAsciiLen(text, size):
76    """Trim a string containing ANSI sequences to the given ASCII length
77
78    The string is trimmed with ANSI sequences being ignored for the length
79    calculation.
80
81    >>> col = Color(COLOR_ALWAYS)
82    >>> text = col.Color(Color.RED, 'abc')
83    >>> len(text)
84    14
85    >>> CalcAsciiLen(TrimAsciiLen(text, 4))
86    3
87    >>> CalcAsciiLen(TrimAsciiLen(text, 2))
88    2
89    >>> text += 'def'
90    >>> CalcAsciiLen(TrimAsciiLen(text, 4))
91    4
92    >>> text += col.Color(Color.RED, 'ghi')
93    >>> CalcAsciiLen(TrimAsciiLen(text, 7))
94    7
95    """
96    if CalcAsciiLen(text) < size:
97        return text
98    pos = 0
99    out = ''
100    left = size
101
102    # Work through each ANSI sequence in turn
103    for m in ansi_escape.finditer(text):
104        # Find the text before the sequence and add it to our string, making
105        # sure it doesn't overflow
106        before = text[pos:m.start()]
107        toadd = before[:left]
108        out += toadd
109
110        # Figure out how much non-ANSI space we have left
111        left -= len(toadd)
112
113        # Add the ANSI sequence and move to the position immediately after it
114        out += m.group()
115        pos = m.start() + len(m.group())
116
117    # Deal with text after the last ANSI sequence
118    after = text[pos:]
119    toadd = after[:left]
120    out += toadd
121
122    return out
123
124
125def Print(text='', newline=True, colour=None, limit_to_line=False):
126    """Handle a line of output to the terminal.
127
128    In test mode this is recorded in a list. Otherwise it is output to the
129    terminal.
130
131    Args:
132        text: Text to print
133        newline: True to add a new line at the end of the text
134        colour: Colour to use for the text
135    """
136    global last_print_len
137
138    if print_test_mode:
139        print_test_list.append(PrintLine(text, newline, colour))
140    else:
141        if colour:
142            col = Color()
143            text = col.Color(colour, text)
144        if newline:
145            print(text)
146            last_print_len = None
147        else:
148            if limit_to_line:
149                cols = shutil.get_terminal_size().columns
150                text = TrimAsciiLen(text, cols)
151            print(text, end='', flush=True)
152            last_print_len = CalcAsciiLen(text)
153
154def PrintClear():
155    """Clear a previously line that was printed with no newline"""
156    global last_print_len
157
158    if last_print_len:
159        print('\r%s\r' % (' '* last_print_len), end='', flush=True)
160        last_print_len = None
161
162def SetPrintTestMode():
163    """Go into test mode, where all printing is recorded"""
164    global print_test_mode
165
166    print_test_mode = True
167
168def GetPrintTestLines():
169    """Get a list of all lines output through Print()
170
171    Returns:
172        A list of PrintLine objects
173    """
174    global print_test_list
175
176    ret = print_test_list
177    print_test_list = []
178    return ret
179
180def EchoPrintTestLines():
181    """Print out the text lines collected"""
182    for line in print_test_list:
183        if line.colour:
184            col = Color()
185            print(col.Color(line.colour, line.text), end='')
186        else:
187            print(line.text, end='')
188        if line.newline:
189            print()
190
191
192class Color(object):
193    """Conditionally wraps text in ANSI color escape sequences."""
194    BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
195    BOLD = -1
196    BRIGHT_START = '\033[1;%dm'
197    NORMAL_START = '\033[22;%dm'
198    BOLD_START = '\033[1m'
199    RESET = '\033[0m'
200
201    def __init__(self, colored=COLOR_IF_TERMINAL):
202        """Create a new Color object, optionally disabling color output.
203
204        Args:
205          enabled: True if color output should be enabled. If False then this
206            class will not add color codes at all.
207        """
208        try:
209            self._enabled = (colored == COLOR_ALWAYS or
210                    (colored == COLOR_IF_TERMINAL and
211                     os.isatty(sys.stdout.fileno())))
212        except:
213            self._enabled = False
214
215    def Start(self, color, bright=True):
216        """Returns a start color code.
217
218        Args:
219          color: Color to use, .e.g BLACK, RED, etc.
220
221        Returns:
222          If color is enabled, returns an ANSI sequence to start the given
223          color, otherwise returns empty string
224        """
225        if self._enabled:
226            base = self.BRIGHT_START if bright else self.NORMAL_START
227            return base % (color + 30)
228        return ''
229
230    def Stop(self):
231        """Returns a stop color code.
232
233        Returns:
234          If color is enabled, returns an ANSI color reset sequence,
235          otherwise returns empty string
236        """
237        if self._enabled:
238            return self.RESET
239        return ''
240
241    def Color(self, color, text, bright=True):
242        """Returns text with conditionally added color escape sequences.
243
244        Keyword arguments:
245          color: Text color -- one of the color constants defined in this
246                  class.
247          text: The text to color.
248
249        Returns:
250          If self._enabled is False, returns the original text. If it's True,
251          returns text with color escape sequences based on the value of
252          color.
253        """
254        if not self._enabled:
255            return text
256        if color == self.BOLD:
257            start = self.BOLD_START
258        else:
259            base = self.BRIGHT_START if bright else self.NORMAL_START
260            start = base % (color + 30)
261        return start + text + self.RESET
262