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