1# Copyright 2016 The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Terminal utilities 16 17This module handles terminal interaction including ANSI color codes. 18""" 19 20import os 21import sys 22 23_path = os.path.realpath(__file__ + '/../..') 24if sys.path[0] != _path: 25 sys.path.insert(0, _path) 26del _path 27 28# pylint: disable=wrong-import-position 29import rh.shell 30 31 32class Color(object): 33 """Conditionally wraps text in ANSI color escape sequences.""" 34 35 BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8) 36 BOLD = -1 37 COLOR_START = '\033[1;%dm' 38 BOLD_START = '\033[1m' 39 RESET = '\033[0m' 40 41 def __init__(self, enabled=None): 42 """Create a new Color object, optionally disabling color output. 43 44 Args: 45 enabled: True if color output should be enabled. If False then this 46 class will not add color codes at all. 47 """ 48 self._enabled = enabled 49 50 def start(self, color): 51 """Returns a start color code. 52 53 Args: 54 color: Color to use, .e.g BLACK, RED, etc. 55 56 Returns: 57 If color is enabled, returns an ANSI sequence to start the given 58 color, otherwise returns empty string 59 """ 60 if self.enabled: 61 return self.COLOR_START % (color + 30) 62 return '' 63 64 def stop(self): 65 """Returns a stop color code. 66 67 Returns: 68 If color is enabled, returns an ANSI color reset sequence, otherwise 69 returns empty string 70 """ 71 if self.enabled: 72 return self.RESET 73 return '' 74 75 def color(self, color, text): 76 """Returns text with conditionally added color escape sequences. 77 78 Args: 79 color: Text color -- one of the color constants defined in this class. 80 text: The text to color. 81 82 Returns: 83 If self._enabled is False, returns the original text. If it's True, 84 returns text with color escape sequences based on the value of color. 85 """ 86 if not self.enabled: 87 return text 88 if color == self.BOLD: 89 start = self.BOLD_START 90 else: 91 start = self.COLOR_START % (color + 30) 92 return start + text + self.RESET 93 94 @property 95 def enabled(self): 96 """See if the colorization is enabled.""" 97 if self._enabled is None: 98 if 'NOCOLOR' in os.environ: 99 self._enabled = not rh.shell.boolean_shell_value( 100 os.environ['NOCOLOR'], False) 101 else: 102 self._enabled = is_tty(sys.stderr) 103 return self._enabled 104 105 106def is_tty(fh): 107 """Returns whether the specified file handle is a TTY. 108 109 Args: 110 fh: File handle to check. 111 112 Returns: 113 True if |fh| is a TTY 114 """ 115 try: 116 return os.isatty(fh.fileno()) 117 except IOError: 118 return False 119 120 121def print_status_line(line, print_newline=False): 122 """Clears the current terminal line, and prints |line|. 123 124 Args: 125 line: String to print. 126 print_newline: Print a newline at the end, if sys.stderr is a TTY. 127 """ 128 if is_tty(sys.stderr): 129 output = '\r' + line + '\x1B[K' 130 if print_newline: 131 output += '\n' 132 else: 133 output = line + '\n' 134 135 sys.stderr.write(output) 136 sys.stderr.flush() 137 138 139def get_input(prompt): 140 """Python 2/3 glue for raw_input/input differences.""" 141 try: 142 # pylint: disable=raw_input-builtin 143 return raw_input(prompt) 144 except NameError: 145 # Python 3 renamed raw_input() to input(), which is safe to call since 146 # it does not evaluate the input. 147 # pylint: disable=bad-builtin,input-builtin 148 return input(prompt) 149 150 151def boolean_prompt(prompt='Do you want to continue?', default=True, 152 true_value='yes', false_value='no', prolog=None): 153 """Helper function for processing boolean choice prompts. 154 155 Args: 156 prompt: The question to present to the user. 157 default: Boolean to return if the user just presses enter. 158 true_value: The text to display that represents a True returned. 159 false_value: The text to display that represents a False returned. 160 prolog: The text to display before prompt. 161 162 Returns: 163 True or False. 164 """ 165 true_value, false_value = true_value.lower(), false_value.lower() 166 true_text, false_text = true_value, false_value 167 if true_value == false_value: 168 raise ValueError('true_value and false_value must differ: got %r' 169 % true_value) 170 171 if default: 172 true_text = true_text[0].upper() + true_text[1:] 173 else: 174 false_text = false_text[0].upper() + false_text[1:] 175 176 prompt = ('\n%s (%s/%s)? ' % (prompt, true_text, false_text)) 177 178 if prolog: 179 prompt = ('\n%s\n%s' % (prolog, prompt)) 180 181 while True: 182 try: 183 response = get_input(prompt).lower() 184 except EOFError: 185 # If the user hits CTRL+D, or stdin is disabled, use the default. 186 print() 187 response = None 188 except KeyboardInterrupt: 189 # If the user hits CTRL+C, just exit the process. 190 print() 191 raise 192 193 if not response: 194 return default 195 if true_value.startswith(response): 196 if not false_value.startswith(response): 197 return True 198 # common prefix between the two... 199 elif false_value.startswith(response): 200 return False 201