summaryrefslogtreecommitdiff
path: root/tools/patman/terminal.py
blob: 60dbce3ce1ffd7da78a4141220b4ce47defd4078 (plain)
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
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
# SPDX-License-Identifier: GPL-2.0+
# Copyright (c) 2011 The Chromium OS Authors.
#

"""Terminal utilities

This module handles terminal interaction including ANSI color codes.
"""

import os
import re
import shutil
import sys

# Selection of when we want our output to be colored
COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)

# Initially, we are set up to print to the terminal
print_test_mode = False
print_test_list = []

# The length of the last line printed without a newline
last_print_len = None

# credit:
# stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python
ansi_escape = re.compile(r'\x1b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')

class PrintLine:
    """A line of text output

    Members:
        text: Text line that was printed
        newline: True to output a newline after the text
        colour: Text colour to use
    """
    def __init__(self, text, newline, colour):
        self.text = text
        self.newline = newline
        self.colour = colour

    def __str__(self):
        return 'newline=%s, colour=%s, text=%s' % (self.newline, self.colour,
                self.text)

def CalcAsciiLen(text):
    """Calculate the length of a string, ignoring any ANSI sequences

    When displayed on a terminal, ANSI sequences don't take any space, so we
    need to ignore them when calculating the length of a string.

    Args:
        text: Text to check

    Returns:
        Length of text, after skipping ANSI sequences

    >>> col = Color(COLOR_ALWAYS)
    >>> text = col.Color(Color.RED, 'abc')
    >>> len(text)
    14
    >>> CalcAsciiLen(text)
    3
    >>>
    >>> text += 'def'
    >>> CalcAsciiLen(text)
    6
    >>> text += col.Color(Color.RED, 'abc')
    >>> CalcAsciiLen(text)
    9
    """
    result = ansi_escape.sub('', text)
    return len(result)

def TrimAsciiLen(text, size):
    """Trim a string containing ANSI sequences to the given ASCII length

    The string is trimmed with ANSI sequences being ignored for the length
    calculation.

    >>> col = Color(COLOR_ALWAYS)
    >>> text = col.Color(Color.RED, 'abc')
    >>> len(text)
    14
    >>> CalcAsciiLen(TrimAsciiLen(text, 4))
    3
    >>> CalcAsciiLen(TrimAsciiLen(text, 2))
    2
    >>> text += 'def'
    >>> CalcAsciiLen(TrimAsciiLen(text, 4))
    4
    >>> text += col.Color(Color.RED, 'ghi')
    >>> CalcAsciiLen(TrimAsciiLen(text, 7))
    7
    """
    if CalcAsciiLen(text) < size:
        return text
    pos = 0
    out = ''
    left = size

    # Work through each ANSI sequence in turn
    for m in ansi_escape.finditer(text):
        # Find the text before the sequence and add it to our string, making
        # sure it doesn't overflow
        before = text[pos:m.start()]
        toadd = before[:left]
        out += toadd

        # Figure out how much non-ANSI space we have left
        left -= len(toadd)

        # Add the ANSI sequence and move to the position immediately after it
        out += m.group()
        pos = m.start() + len(m.group())

    # Deal with text after the last ANSI sequence
    after = text[pos:]
    toadd = after[:left]
    out += toadd

    return out


def Print(text='', newline=True, colour=None, limit_to_line=False, bright=True):
    """Handle a line of output to the terminal.

    In test mode this is recorded in a list. Otherwise it is output to the
    terminal.

    Args:
        text: Text to print
        newline: True to add a new line at the end of the text
        colour: Colour to use for the text
    """
    global last_print_len

    if print_test_mode:
        print_test_list.append(PrintLine(text, newline, colour))
    else:
        if colour:
            col = Color()
            text = col.Color(colour, text, bright=bright)
        if newline:
            print(text)
            last_print_len = None
        else:
            if limit_to_line:
                cols = shutil.get_terminal_size().columns
                text = TrimAsciiLen(text, cols)
            print(text, end='', flush=True)
            last_print_len = CalcAsciiLen(text)

def PrintClear():
    """Clear a previously line that was printed with no newline"""
    global last_print_len

    if last_print_len:
        print('\r%s\r' % (' '* last_print_len), end='', flush=True)
        last_print_len = None

def SetPrintTestMode():
    """Go into test mode, where all printing is recorded"""
    global print_test_mode

    print_test_mode = True

def GetPrintTestLines():
    """Get a list of all lines output through Print()

    Returns:
        A list of PrintLine objects
    """
    global print_test_list

    ret = print_test_list
    print_test_list = []
    return ret

def EchoPrintTestLines():
    """Print out the text lines collected"""
    for line in print_test_list:
        if line.colour:
            col = Color()
            print(col.Color(line.colour, line.text), end='')
        else:
            print(line.text, end='')
        if line.newline:
            print()


class Color(object):
    """Conditionally wraps text in ANSI color escape sequences."""
    BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
    BOLD = -1
    BRIGHT_START = '\033[1;%dm'
    NORMAL_START = '\033[22;%dm'
    BOLD_START = '\033[1m'
    RESET = '\033[0m'

    def __init__(self, colored=COLOR_IF_TERMINAL):
        """Create a new Color object, optionally disabling color output.

        Args:
          enabled: True if color output should be enabled. If False then this
            class will not add color codes at all.
        """
        try:
            self._enabled = (colored == COLOR_ALWAYS or
                    (colored == COLOR_IF_TERMINAL and
                     os.isatty(sys.stdout.fileno())))
        except:
            self._enabled = False

    def Start(self, color, bright=True):
        """Returns a start color code.

        Args:
          color: Color to use, .e.g BLACK, RED, etc.

        Returns:
          If color is enabled, returns an ANSI sequence to start the given
          color, otherwise returns empty string
        """
        if self._enabled:
            base = self.BRIGHT_START if bright else self.NORMAL_START
            return base % (color + 30)
        return ''

    def Stop(self):
        """Returns a stop color code.

        Returns:
          If color is enabled, returns an ANSI color reset sequence,
          otherwise returns empty string
        """
        if self._enabled:
            return self.RESET
        return ''

    def Color(self, color, text, bright=True):
        """Returns text with conditionally added color escape sequences.

        Keyword arguments:
          color: Text color -- one of the color constants defined in this
                  class.
          text: The text to color.

        Returns:
          If self._enabled is False, returns the original text. If it's True,
          returns text with color escape sequences based on the value of
          color.
        """
        if not self._enabled:
            return text
        if color == self.BOLD:
            start = self.BOLD_START
        else:
            base = self.BRIGHT_START if bright else self.NORMAL_START
            start = base % (color + 30)
        return start + text + self.RESET