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 class Color(object):
  18     """Conditionally wraps text in ANSI color escape sequences."""
  19     BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
  20     BOLD = -1
  21     BRIGHT_START = '\033[1;%dm'
  22     NORMAL_START = '\033[22;%dm'
  23     BOLD_START = '\033[1m'
  24     RESET = '\033[0m'
  25
  26     def __init__(self, colored=COLOR_IF_TERMINAL):
  27         """Create a new Color object, optionally disabling color output.
  28
  29         Args:
  30           enabled: True if color output should be enabled. If False then this
  31             class will not add color codes at all.
  32         """
  33         try:
  34             self._enabled = (colored == COLOR_ALWAYS or
  35                     (colored == COLOR_IF_TERMINAL and
  36                      os.isatty(sys.stdout.fileno())))
  37         except:
  38             self._enabled = False
  39
  40     def Start(self, color, bright=True):
  41         """Returns a start color code.
  42
  43         Args:
  44           color: Color to use, .e.g BLACK, RED, etc.
  45
  46         Returns:
  47           If color is enabled, returns an ANSI sequence to start the given
  48           color, otherwise returns empty string
  49         """
  50         if self._enabled:
  51             base = self.BRIGHT_START if bright else self.NORMAL_START
  52             return base % (color + 30)
  53         return ''
  54
  55     def Stop(self):
  56         """Retruns a stop color code.
  57
  58         Returns:
  59           If color is enabled, returns an ANSI color reset sequence,
  60           otherwise returns empty string
  61         """
  62         if self._enabled:
  63             return self.RESET
  64         return ''
  65
  66     def Color(self, color, text, bright=True):
  67         """Returns text with conditionally added color escape sequences.
  68
  69         Keyword arguments:
  70           color: Text color -- one of the color constants defined in this
  71                   class.
  72           text: The text to color.
  73
  74         Returns:
  75           If self._enabled is False, returns the original text. If it's True,
  76           returns text with color escape sequences based on the value of
  77           color.
  78         """
  79         if not self._enabled:
  80             return text
  81         if color == self.BOLD:
  82             start = self.BOLD_START
  83         else:
  84             base = self.BRIGHT_START if bright else self.NORMAL_START
  85             start = base % (color + 30)
  86         return start + text + self.RESET