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
22from typing import List, Optional
23
24_path = os.path.realpath(__file__ + '/../..')
25if sys.path[0] != _path:
26    sys.path.insert(0, _path)
27del _path
28
29# pylint: disable=wrong-import-position
30import rh.shell
31
32
33# This will erase all content in the current line after the cursor.  This is
34# useful for partial updates & progress messages as the terminal can display
35# it better.
36CSI_ERASE_LINE_AFTER = '\x1b[K'
37
38
39class Color(object):
40    """Conditionally wraps text in ANSI color escape sequences."""
41
42    BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
43    BOLD = -1
44    COLOR_START = '\033[1;%dm'
45    BOLD_START = '\033[1m'
46    RESET = '\033[m'
47
48    def __init__(self, enabled=None):
49        """Create a new Color object, optionally disabling color output.
50
51        Args:
52          enabled: True if color output should be enabled.  If False then this
53              class will not add color codes at all.
54        """
55        self._enabled = enabled
56
57    def start(self, color):
58        """Returns a start color code.
59
60        Args:
61          color: Color to use, e.g. BLACK, RED, etc...
62
63        Returns:
64          If color is enabled, returns an ANSI sequence to start the given
65          color, otherwise returns empty string
66        """
67        if self.enabled:
68            return self.COLOR_START % (color + 30)
69        return ''
70
71    def stop(self):
72        """Returns a stop color code.
73
74        Returns:
75          If color is enabled, returns an ANSI color reset sequence, otherwise
76          returns empty string
77        """
78        if self.enabled:
79            return self.RESET
80        return ''
81
82    def color(self, color, text):
83        """Returns text with conditionally added color escape sequences.
84
85        Args:
86          color: Text color -- one of the color constants defined in this class.
87          text: The text to color.
88
89        Returns:
90          If self._enabled is False, returns the original text.  If it's True,
91          returns text with color escape sequences based on the value of color.
92        """
93        if not self.enabled:
94            return text
95        if color == self.BOLD:
96            start = self.BOLD_START
97        else:
98            start = self.COLOR_START % (color + 30)
99        return start + text + self.RESET
100
101    @property
102    def enabled(self):
103        """See if the colorization is enabled."""
104        if self._enabled is None:
105            if 'NOCOLOR' in os.environ:
106                self._enabled = not rh.shell.boolean_shell_value(
107                    os.environ['NOCOLOR'], False)
108            else:
109                self._enabled = sys.stderr.isatty()
110        return self._enabled
111
112
113def print_status_line(line, print_newline=False):
114    """Clears the current terminal line, and prints |line|.
115
116    Args:
117      line: String to print.
118      print_newline: Print a newline at the end, if sys.stderr is a TTY.
119    """
120    if sys.stderr.isatty():
121        output = '\r' + line + CSI_ERASE_LINE_AFTER
122        if print_newline:
123            output += '\n'
124    else:
125        output = line + '\n'
126
127    sys.stderr.write(output)
128    sys.stderr.flush()
129
130
131def str_prompt(
132    prompt: str,
133    choices: List[str],
134    lower: bool = True,
135) -> Optional[str]:
136    """Helper function for processing user input.
137
138    Args:
139        prompt: The question to present to the user.
140        lower: Whether to lowercase the response.
141
142    Returns:
143        The string the user entered, or None if EOF (e.g. Ctrl+D).
144    """
145    prompt = f'{prompt} ({"/".join(choices)})? '
146    try:
147        result = input(prompt)
148        return result.lower() if lower else result
149    except EOFError:
150        # If the user hits Ctrl+D, or stdin is disabled, use the default.
151        print()
152        return None
153    except KeyboardInterrupt:
154        # If the user hits Ctrl+C, just exit the process.
155        print()
156        raise
157
158
159def boolean_prompt(prompt='Do you want to continue?', default=True,
160                   true_value='yes', false_value='no', prolog=None):
161    """Helper function for processing boolean choice prompts.
162
163    Args:
164      prompt: The question to present to the user.
165      default: Boolean to return if the user just presses enter.
166      true_value: The text to display that represents a True returned.
167      false_value: The text to display that represents a False returned.
168      prolog: The text to display before prompt.
169
170    Returns:
171      True or False.
172    """
173    true_value, false_value = true_value.lower(), false_value.lower()
174    true_text, false_text = true_value, false_value
175    if true_value == false_value:
176        raise ValueError(
177            f'true_value and false_value must differ: got {true_value!r}')
178
179    if default:
180        true_text = true_text[0].upper() + true_text[1:]
181    else:
182        false_text = false_text[0].upper() + false_text[1:]
183
184    if prolog:
185        prompt = f'\n{prolog}\n{prompt}'
186    prompt = '\n' + prompt
187
188    while True:
189        response = str_prompt(prompt, choices=(true_text, false_text))
190        if not response:
191            return default
192        if true_value.startswith(response):
193            if not false_value.startswith(response):
194                return True
195            # common prefix between the two...
196        elif false_value.startswith(response):
197            return False
198