• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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