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 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 class PrintLine:
  23     """A line of text output
  24
  25     Members:
  26         text: Text line that was printed
  27         newline: True to output a newline after the text
  28         colour: Text colour to use
  29     """
  30     def __init__(self, text, newline, colour):
  31         self.text = text
  32         self.newline = newline
  33         self.colour = colour
  34
  35     def __str__(self):
  36         return 'newline=%s, colour=%s, text=%s' % (self.newline, self.colour,
  37                 self.text)
  38
  39 def Print(text='', newline=True, colour=None):
  40     """Handle a line of output to the terminal.
  41
  42     In test mode this is recorded in a list. Otherwise it is output to the
  43     terminal.
  44
  45     Args:
  46         text: Text to print
  47         newline: True to add a new line at the end of the text
  48         colour: Colour to use for the text
  49     """
  50     if print_test_mode:
  51         print_test_list.append(PrintLine(text, newline, colour))
  52     else:
  53         if colour:
  54             col = Color()
  55             text = col.Color(colour, text)
  56         print(text, end='')
  57         if newline:
  58             print()
  59         else:
  60             sys.stdout.flush()
  61
  62 def SetPrintTestMode():
  63     """Go into test mode, where all printing is recorded"""
  64     global print_test_mode
  65
  66     print_test_mode = True
  67
  68 def GetPrintTestLines():
  69     """Get a list of all lines output through Print()
  70
  71     Returns:
  72         A list of PrintLine objects
  73     """
  74     global print_test_list
  75
  76     ret = print_test_list
  77     print_test_list = []
  78     return ret
  79
  80 def EchoPrintTestLines():
  81     """Print out the text lines collected"""
  82     for line in print_test_list:
  83         if line.colour:
  84             col = Color()
  85             print(col.Color(line.colour, line.text), end='')
  86         else:
  87             print(line.text, end='')
  88         if line.newline:
  89             print()
  90
  91
  92 class Color(object):
  93     """Conditionally wraps text in ANSI color escape sequences."""
  94     BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
  95     BOLD = -1
  96     BRIGHT_START = '\033[1;%dm'
  97     NORMAL_START = '\033[22;%dm'
  98     BOLD_START = '\033[1m'
  99     RESET = '\033[0m'
 100
 101     def __init__(self, colored=COLOR_IF_TERMINAL):
 102         """Create a new Color object, optionally disabling color output.
 103
 104         Args:
 105           enabled: True if color output should be enabled. If False then this
 106             class will not add color codes at all.
 107         """
 108         try:
 109             self._enabled = (colored == COLOR_ALWAYS or
 110                     (colored == COLOR_IF_TERMINAL and
 111                      os.isatty(sys.stdout.fileno())))
 112         except:
 113             self._enabled = False
 114
 115     def Start(self, color, bright=True):
 116         """Returns a start color code.
 117
 118         Args:
 119           color: Color to use, .e.g BLACK, RED, etc.
 120
 121         Returns:
 122           If color is enabled, returns an ANSI sequence to start the given
 123           color, otherwise returns empty string
 124         """
 125         if self._enabled:
 126             base = self.BRIGHT_START if bright else self.NORMAL_START
 127             return base % (color + 30)
 128         return ''
 129
 130     def Stop(self):
 131         """Retruns a stop color code.
 132
 133         Returns:
 134           If color is enabled, returns an ANSI color reset sequence,
 135           otherwise returns empty string
 136         """
 137         if self._enabled:
 138             return self.RESET
 139         return ''
 140
 141     def Color(self, color, text, bright=True):
 142         """Returns text with conditionally added color escape sequences.
 143
 144         Keyword arguments:
 145           color: Text color -- one of the color constants defined in this
 146                   class.
 147           text: The text to color.
 148
 149         Returns:
 150           If self._enabled is False, returns the original text. If it's True,
 151           returns text with color escape sequences based on the value of
 152           color.
 153         """
 154         if not self._enabled:
 155             return text
 156         if color == self.BOLD:
 157             start = self.BOLD_START
 158         else:
 159             base = self.BRIGHT_START if bright else self.NORMAL_START
 160             start = base % (color + 30)
 161         return start + text + self.RESET