1 files changed, 42 insertions, 48 deletions

diff --git a/rh/terminal.py b/rh/terminal.py
index c549f12..a6f31d9 100644
--- a/rh/terminal.py
+++ b/rh/terminal.py
@@ -1,4 +1,3 @@
-# -*- coding:utf-8 -*-
 # Copyright 2016 The Android Open Source Project
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,10 +17,9 @@
 This module handles terminal interaction including ANSI color codes.
 """
 
-from __future__ import print_function
-
 import os
 import sys
+from typing import List, Optional
 
 _path = os.path.realpath(__file__ + '/../..')
 if sys.path[0] != _path:
@@ -32,6 +30,12 @@ del _path
 import rh.shell
 
 
+# This will erase all content in the current line after the cursor.  This is
+# useful for partial updates & progress messages as the terminal can display
+# it better.
+CSI_ERASE_LINE_AFTER = '\x1b[K'
+
+
 class Color(object):
     """Conditionally wraps text in ANSI color escape sequences."""
 
@@ -39,7 +43,7 @@ class Color(object):
     BOLD = -1
     COLOR_START = '\033[1;%dm'
     BOLD_START = '\033[1m'
-    RESET = '\033[0m'
+    RESET = '\033[m'
 
     def __init__(self, enabled=None):
         """Create a new Color object, optionally disabling color output.
@@ -54,7 +58,7 @@ class Color(object):
         """Returns a start color code.
 
         Args:
-          color: Color to use, .e.g BLACK, RED, etc.
+          color: Color to use, e.g. BLACK, RED, etc...
 
         Returns:
           If color is enabled, returns an ANSI sequence to start the given
@@ -102,25 +106,10 @@ class Color(object):
                 self._enabled = not rh.shell.boolean_shell_value(
                     os.environ['NOCOLOR'], False)
             else:
-                self._enabled = is_tty(sys.stderr)
+                self._enabled = sys.stderr.isatty()
         return self._enabled
 
 
-def is_tty(fh):
-    """Returns whether the specified file handle is a TTY.
-
-    Args:
-      fh: File handle to check.
-
-    Returns:
-      True if |fh| is a TTY
-    """
-    try:
-        return os.isatty(fh.fileno())
-    except IOError:
-        return False
-
-
 def print_status_line(line, print_newline=False):
     """Clears the current terminal line, and prints |line|.
 
@@ -128,8 +117,8 @@ def print_status_line(line, print_newline=False):
       line: String to print.
       print_newline: Print a newline at the end, if sys.stderr is a TTY.
     """
-    if is_tty(sys.stderr):
-        output = '\r' + line + '\x1B[K'
+    if sys.stderr.isatty():
+        output = '\r' + line + CSI_ERASE_LINE_AFTER
         if print_newline:
             output += '\n'
     else:
@@ -139,16 +128,32 @@ def print_status_line(line, print_newline=False):
     sys.stderr.flush()
 
 
-def get_input(prompt):
-    """Python 2/3 glue for raw_input/input differences."""
+def str_prompt(
+    prompt: str,
+    choices: List[str],
+    lower: bool = True,
+) -> Optional[str]:
+    """Helper function for processing user input.
+
+    Args:
+        prompt: The question to present to the user.
+        lower: Whether to lowercase the response.
+
+    Returns:
+        The string the user entered, or None if EOF (e.g. Ctrl+D).
+    """
+    prompt = f'{prompt} ({"/".join(choices)})? '
     try:
-        # pylint: disable=raw_input-builtin
-        return raw_input(prompt)
-    except NameError:
-        # Python 3 renamed raw_input() to input(), which is safe to call since
-        # it does not evaluate the input.
-        # pylint: disable=bad-builtin,input-builtin
-        return input(prompt)
+        result = input(prompt)
+        return result.lower() if lower else result
+    except EOFError:
+        # If the user hits Ctrl+D, or stdin is disabled, use the default.
+        print()
+        return None
+    except KeyboardInterrupt:
+        # If the user hits Ctrl+C, just exit the process.
+        print()
+        raise
 
 
 def boolean_prompt(prompt='Do you want to continue?', default=True,
@@ -168,31 +173,20 @@ def boolean_prompt(prompt='Do you want to continue?', default=True,
     true_value, false_value = true_value.lower(), false_value.lower()
     true_text, false_text = true_value, false_value
     if true_value == false_value:
-        raise ValueError('true_value and false_value must differ: got %r'
-                         % true_value)
+        raise ValueError(
+            f'true_value and false_value must differ: got {true_value!r}')
 
     if default:
         true_text = true_text[0].upper() + true_text[1:]
     else:
         false_text = false_text[0].upper() + false_text[1:]
 
-    prompt = ('\n%s (%s/%s)? ' % (prompt, true_text, false_text))
-
     if prolog:
-        prompt = ('\n%s\n%s' % (prolog, prompt))
+        prompt = f'\n{prolog}\n{prompt}'
+    prompt = '\n' + prompt
 
     while True:
-        try:
-            response = get_input(prompt).lower()
-        except EOFError:
-            # If the user hits CTRL+D, or stdin is disabled, use the default.
-            print()
-            response = None
-        except KeyboardInterrupt:
-            # If the user hits CTRL+C, just exit the process.
-            print()
-            raise
-
+        response = str_prompt(prompt, choices=(true_text, false_text))
         if not response:
             return default
         if true_value.startswith(response):