tools/patman/terminal.py

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