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