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