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         self._enabled = (colored == COLOR_ALWAYS or
  34             (colored == COLOR_IF_TERMINAL and os.isatty(sys.stdout.fileno())))
  35
  36     def Start(self, color, bright=True):
  37         """Returns a start color code.
  38
  39         Args:
  40           color: Color to use, .e.g BLACK, RED, etc.
  41
  42         Returns:
  43           If color is enabled, returns an ANSI sequence to start the given
  44           color, otherwise returns empty string
  45         """
  46         if self._enabled:
  47             base = self.BRIGHT_START if bright else self.NORMAL_START
  48             return base % (color + 30)
  49         return ''
  50
  51     def Stop(self):
  52         """Retruns a stop color code.
  53
  54         Returns:
  55           If color is enabled, returns an ANSI color reset sequence,
  56           otherwise returns empty string
  57         """
  58         if self._enabled:
  59             return self.RESET
  60         return ''
  61
  62     def Color(self, color, text, bright=True):
  63         """Returns text with conditionally added color escape sequences.
  64
  65         Keyword arguments:
  66           color: Text color -- one of the color constants defined in this
  67                   class.
  68           text: The text to color.
  69
  70         Returns:
  71           If self._enabled is False, returns the original text. If it's True,
  72           returns text with color escape sequences based on the value of
  73           color.
  74         """
  75         if not self._enabled:
  76             return text
  77         if color == self.BOLD:
  78             start = self.BOLD_START
  79         else:
  80             base = self.BRIGHT_START if bright else self.NORMAL_START
  81             start = base % (color + 30)
  82         return start + text + self.RESET