1# Copyright (c) 2011 The Chromium OS Authors. 2# 3# See file CREDITS for list of people who contributed to this 4# project. 5# 6# This program is free software; you can redistribute it and/or 7# modify it under the terms of the GNU General Public License as 8# published by the Free Software Foundation; either version 2 of 9# the License, or (at your option) any later version. 10# 11# This program is distributed in the hope that it will be useful, 12# but WITHOUT ANY WARRANTY; without even the implied warranty of 13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14# GNU General Public License for more details. 15# 16# You should have received a copy of the GNU General Public License 17# along with this program; if not, write to the Free Software 18# Foundation, Inc., 59 Temple Place, Suite 330, Boston, 19# MA 02111-1307 USA 20# 21 22"""Terminal utilities 23 24This module handles terminal interaction including ANSI color codes. 25""" 26 27import os 28import sys 29 30# Selection of when we want our output to be colored 31COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3) 32 33class Color(object): 34 """Conditionally wraps text in ANSI color escape sequences.""" 35 BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8) 36 BOLD = -1 37 BRIGHT_START = '\033[1;%dm' 38 NORMAL_START = '\033[22;%dm' 39 BOLD_START = '\033[1m' 40 RESET = '\033[0m' 41 42 def __init__(self, colored=COLOR_IF_TERMINAL): 43 """Create a new Color object, optionally disabling color output. 44 45 Args: 46 enabled: True if color output should be enabled. If False then this 47 class will not add color codes at all. 48 """ 49 self._enabled = (colored == COLOR_ALWAYS or 50 (colored == COLOR_IF_TERMINAL and os.isatty(sys.stdout.fileno()))) 51 52 def Start(self, color, bright=True): 53 """Returns a start color code. 54 55 Args: 56 color: Color to use, .e.g BLACK, RED, etc. 57 58 Returns: 59 If color is enabled, returns an ANSI sequence to start the given color, 60 otherwise returns empty string 61 """ 62 if self._enabled: 63 base = self.BRIGHT_START if bright else self.NORMAL_START 64 return base % (color + 30) 65 return '' 66 67 def Stop(self): 68 """Retruns a stop color code. 69 70 Returns: 71 If color is enabled, returns an ANSI color reset sequence, otherwise 72 returns empty string 73 """ 74 if self._enabled: 75 return self.RESET 76 return '' 77 78 def Color(self, color, text, bright=True): 79 """Returns text with conditionally added color escape sequences. 80 81 Keyword arguments: 82 color: Text color -- one of the color constants defined in this class. 83 text: The text to color. 84 85 Returns: 86 If self._enabled is False, returns the original text. If it's True, 87 returns text with color escape sequences based on the value of color. 88 """ 89 if not self._enabled: 90 return text 91 if color == self.BOLD: 92 start = self.BOLD_START 93 else: 94 base = self.BRIGHT_START if bright else self.NORMAL_START 95 start = base % (color + 30) 96 return start + text + self.RESET 97