Hammer on that run_tests.py thing again.

[python_utils.git] / input_utils.py

diff --git a/input_utils.py b/input_utils.py
index 635e349009f174c9d7a69b68b38a93c4abf5f4bc..0b32eea6f55c5d2fe01466bda0d21ae04c7b13f3 100644 (file)
--- a/input_utils.py
+++ b/input_utils.py
@@ -22,8 +22,24 @@ def single_keystroke_response(
     prompt: str = None,
     default_response: str = None,
     timeout_seconds: int = None,
-) -> str:
-    """Get a single keystroke response to a prompt."""
+) -> Optional[str]:  # None if timeout w/o keystroke
+    """Get a single keystroke response to a prompt and returns it.
+
+    Args:
+        valid_responses: a list of strings that are considered to be
+            valid keystrokes to be accepted.  If None, we accept
+            anything.
+        prompt: the prompt to print before watching keystrokes.  If
+            None, skip this.
+        default_response: the response to return if the timeout
+            expires.  If None, skip this.
+        timeout_seconds: number of seconds to wait before timing out
+            and returning the default_response.  If None, wait forever.
+
+    Returns:
+        The keystroke the user pressed.  If the user pressed a special
+        keystroke like ^C or ^Z, we raise a KeyboardInterrupt exception.
+    """
 
     def _handle_timeout(signum, frame) -> None:
         raise exceptions.TimeoutError()
@@ -49,6 +65,7 @@ def single_keystroke_response(
             if timeout_seconds is not None:
                 signal.alarm(0)
 
+    response = None
     if prompt is not None:
         print(prompt, end="")
         sys.stdout.flush()
@@ -60,23 +77,60 @@ def single_keystroke_response(
     except exceptions.TimeoutError:
         if default_response is not None:
             response = default_response
-    if prompt is not None:
+    if prompt and response:
         print(response)
     return response
 
 
-def yn_response(prompt: str = None, *, timeout_seconds=None) -> str:
-    """Get a Y/N response to a prompt."""
+def yn_response(prompt: str = None, *, timeout_seconds=None) -> Optional[str]:
+    """Get a Y/N response to a prompt.
+
+    Args:
+        prompt: the user prompt or None to skip this
+        timeout_seconds: the number of seconds to wait for a response or
+            None to wait forever.
 
-    return single_keystroke_response(
+    Returns:
+        A lower case 'y' or 'n'.  Or None if the timeout expires with
+        no input from the user.  Or raises a KeyboardInterrupt if the
+        user pressed a special key such as ^C or ^Z.
+    """
+    yn = single_keystroke_response(
         ["y", "n", "Y", "N"], prompt=prompt, timeout_seconds=timeout_seconds
-    ).lower()
+    )
+    if yn:
+        yn = yn.lower()
+    return yn
+
 
+def press_any_key(
+    prompt: str = "Press any key to continue...", *, timeout_seconds=None
+) -> Optional[str]:
+    """Press any key to continue..."""
 
-def press_any_key(prompt: str = "Press any key to continue...", *, timeout_seconds=None) -> str:
     return single_keystroke_response(None, prompt=prompt, timeout_seconds=timeout_seconds)
 
 
+def up_down_enter() -> Optional[str]:
+    """Respond to UP, DOWN or ENTER events for simple menus without
+    the need for curses."""
+
+    os_special_keystrokes = [3, 26]  # ^C, ^Z
+    while True:
+        key = readchar.readkey()
+        if len(key) == 1:
+            if ord(key) in os_special_keystrokes:
+                return None
+            if ord(key) == 13:
+                return 'enter'
+        elif len(key) == 3:
+            if ord(key[0]) == 27 and ord(key[1]) == 91:
+                if ord(key[2]) == 65:
+                    return "up"
+                elif ord(key[2]) == 66:
+                    return "down"
+
+
 def keystroke_helper() -> None:
     """Misc util to watch keystrokes and report what they were."""