src/third_party/chromite/lib/terminal.py

   1 # Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Terminal utilities
   6
   7 This module handles terminal interaction including ANSI color codes.
   8 """
   9
  10 import os
  11
  12 from chromite.lib import cros_build_lib
  13
  14
  15 class Color(object):
  16   """Conditionally wraps text in ANSI color escape sequences."""
  17   BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
  18   BOLD = -1
  19   COLOR_START = '\033[1;%dm'
  20   BOLD_START = '\033[1m'
  21   RESET = '\033[0m'
  22
  23   def __init__(self, enabled=None):
  24     """Create a new Color object, optionally disabling color output.
  25
  26     Args:
  27       enabled: True if color output should be enabled. If False then this
  28         class will not add color codes at all.
  29     """
  30     self._enabled = self.UserEnabled() if enabled is None else enabled
  31
  32   def Start(self, color):
  33     """Returns a start color code.
  34
  35     Args:
  36       color: Color to use, .e.g BLACK, RED, etc.
  37
  38     Returns:
  39       If color is enabled, returns an ANSI sequence to start the given color,
  40       otherwise returns empty string
  41     """
  42     if self._enabled:
  43       return self.COLOR_START % (color + 30)
  44     return ''
  45
  46   def Stop(self):
  47     """Retruns a stop color code.
  48
  49     Returns:
  50       If color is enabled, returns an ANSI color reset sequence, otherwise
  51       returns empty string
  52     """
  53     if self._enabled:
  54       return self.RESET
  55     return ''
  56
  57   def Color(self, color, text):
  58     """Returns text with conditionally added color escape sequences.
  59
  60     Keyword arguments:
  61       color: Text color -- one of the color constants defined in this class.
  62       text: The text to color.
  63
  64     Returns:
  65       If self._enabled is False, returns the original text. If it's True,
  66       returns text with color escape sequences based on the value of color.
  67     """
  68     if not self._enabled:
  69       return text
  70     if color == self.BOLD:
  71       start = self.BOLD_START
  72     else:
  73       start = self.COLOR_START % (color + 30)
  74     return start + text + self.RESET
  75
  76   @staticmethod
  77   def UserEnabled():
  78     """See if the global colorization preference is enabled ($NOCOLOR env)"""
  79     return not cros_build_lib.BooleanShellValue(
  80         os.environ.get('NOCOLOR'), msg='$NOCOLOR env var is invalid',
  81         default=False)