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