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