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 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 87