1# Copyright (c) 2011 The Chromium OS Authors. 2# 3# SPDX-License-Identifier: GPL-2.0+ 4# 5 6"""Terminal utilities 7 8This module handles terminal interaction including ANSI color codes. 9""" 10 11import os 12import sys 13 14# Selection of when we want our output to be colored 15COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3) 16 17class 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 color, 44 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, otherwise 56 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 class. 67 text: The text to color. 68 69 Returns: 70 If self._enabled is False, returns the original text. If it's True, 71 returns text with color escape sequences based on the value of color. 72 """ 73 if not self._enabled: 74 return text 75 if color == self.BOLD: 76 start = self.BOLD_START 77 else: 78 base = self.BRIGHT_START if bright else self.NORMAL_START 79 start = base % (color + 30) 80 return start + text + self.RESET 81