tools/patman/terminal.py

   1 # SPDX-License-Identifier: GPL-2.0+
   2 # Copyright (c) 2011 The Chromium OS Authors.
   3 #
   4
   5 """Terminal utilities
   6
   7 This module handles terminal interaction including ANSI color codes.
   8 """
   9
  10 import os
  11 import re
  12 import shutil
  13 import sys
  14
  15 # Selection of when we want our output to be colored
  16 COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)
  17
  18 # Initially, we are set up to print to the terminal
  19 print_test_mode = False
  20 print_test_list = []
  21
  22 # The length of the last line printed without a newline
  23 last_print_len = None
  24
  25 # credit:
  26 # stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python
  27 ansi_escape = re.compile(r'\x1b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
  28
  29 class 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
  46 def 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
  75 def 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
 125 def 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
 154 def 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
 162 def SetPrintTestMode():
 163     """Go into test mode, where all printing is recorded"""
 164     global print_test_mode
 165
 166     print_test_mode = True
 167
 168 def 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
 180 def 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
 192 class 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