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