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 17# Initially, we are set up to print to the terminal 18print_test_mode = False 19print_test_list = [] 20 21class PrintLine: 22 """A line of text output 23 24 Members: 25 text: Text line that was printed 26 newline: True to output a newline after the text 27 colour: Text colour to use 28 """ 29 def __init__(self, text, newline, colour): 30 self.text = text 31 self.newline = newline 32 self.colour = colour 33 34 def __str__(self): 35 return 'newline=%s, colour=%s, text=%s' % (self.newline, self.colour, 36 self.text) 37 38def Print(text='', newline=True, colour=None): 39 """Handle a line of output to the terminal. 40 41 In test mode this is recorded in a list. Otherwise it is output to the 42 terminal. 43 44 Args: 45 text: Text to print 46 newline: True to add a new line at the end of the text 47 colour: Colour to use for the text 48 """ 49 if print_test_mode: 50 print_test_list.append(PrintLine(text, newline, colour)) 51 else: 52 if colour: 53 col = Color() 54 text = col.Color(colour, text) 55 print text, 56 if newline: 57 print 58 59def SetPrintTestMode(): 60 """Go into test mode, where all printing is recorded""" 61 global print_test_mode 62 63 print_test_mode = True 64 65def GetPrintTestLines(): 66 """Get a list of all lines output through Print() 67 68 Returns: 69 A list of PrintLine objects 70 """ 71 global print_test_list 72 73 ret = print_test_list 74 print_test_list = [] 75 return ret 76 77def EchoPrintTestLines(): 78 """Print out the text lines collected""" 79 for line in print_test_list: 80 if line.colour: 81 col = Color() 82 print col.Color(line.colour, line.text), 83 else: 84 print line.text, 85 if line.newline: 86 print 87 88 89class Color(object): 90 """Conditionally wraps text in ANSI color escape sequences.""" 91 BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8) 92 BOLD = -1 93 BRIGHT_START = '\033[1;%dm' 94 NORMAL_START = '\033[22;%dm' 95 BOLD_START = '\033[1m' 96 RESET = '\033[0m' 97 98 def __init__(self, colored=COLOR_IF_TERMINAL): 99 """Create a new Color object, optionally disabling color output. 100 101 Args: 102 enabled: True if color output should be enabled. If False then this 103 class will not add color codes at all. 104 """ 105 try: 106 self._enabled = (colored == COLOR_ALWAYS or 107 (colored == COLOR_IF_TERMINAL and 108 os.isatty(sys.stdout.fileno()))) 109 except: 110 self._enabled = False 111 112 def Start(self, color, bright=True): 113 """Returns a start color code. 114 115 Args: 116 color: Color to use, .e.g BLACK, RED, etc. 117 118 Returns: 119 If color is enabled, returns an ANSI sequence to start the given 120 color, otherwise returns empty string 121 """ 122 if self._enabled: 123 base = self.BRIGHT_START if bright else self.NORMAL_START 124 return base % (color + 30) 125 return '' 126 127 def Stop(self): 128 """Retruns a stop color code. 129 130 Returns: 131 If color is enabled, returns an ANSI color reset sequence, 132 otherwise returns empty string 133 """ 134 if self._enabled: 135 return self.RESET 136 return '' 137 138 def Color(self, color, text, bright=True): 139 """Returns text with conditionally added color escape sequences. 140 141 Keyword arguments: 142 color: Text color -- one of the color constants defined in this 143 class. 144 text: The text to color. 145 146 Returns: 147 If self._enabled is False, returns the original text. If it's True, 148 returns text with color escape sequences based on the value of 149 color. 150 """ 151 if not self._enabled: 152 return text 153 if color == self.BOLD: 154 start = self.BOLD_START 155 else: 156 base = self.BRIGHT_START if bright else self.NORMAL_START 157 start = base % (color + 30) 158 return start + text + self.RESET 159