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