Fix bug with timeouts.

[python_utils.git] / input_utils.py

diff --git a/input_utils.py b/input_utils.py
index e0b457d5a0e73ffc43d3d83ff09228328e6a641e..5e36db1ba853884636a2df7a5007004aa086760b 100644 (file)
--- a/input_utils.py
+++ b/input_utils.py
@@ -1,34 +1,35 @@
 #!/usr/bin/env python3
 
+# © Copyright 2021-2022, Scott Gasch
+
 """Utilities related to user input."""
 
 import logging
 import signal
 import sys
-from typing import List
+from typing import List, Optional
 
 import readchar  # type: ignore
 
 import exceptions
 
-
 logger = logging.getLogger(__file__)
 
 
 def single_keystroke_response(
-    valid_responses: List[str],
+    valid_responses: Optional[List[str]],  # None = accept anything
     *,
     prompt: str = None,
     default_response: str = None,
     timeout_seconds: int = None,
-) -> str:
+) -> Optional[str]:  # None if timeout w/o keystroke
     """Get a single keystroke response to a prompt."""
 
     def _handle_timeout(signum, frame) -> None:
         raise exceptions.TimeoutError()
 
     def _single_keystroke_response_internal(
-        valid_responses: List[str], timeout_seconds=None
+        valid_responses: Optional[List[str]], timeout_seconds: int = None
     ) -> str:
         os_special_keystrokes = [3, 26]  # ^C, ^Z
         if timeout_seconds is not None:
@@ -38,8 +39,8 @@ def single_keystroke_response(
         try:
             while True:
                 response = readchar.readchar()
-                logger.debug(f'Keystroke: {ord(response)}')
-                if response in valid_responses:
+                logger.debug('Keystroke: 0x%x', ord(response))
+                if not valid_responses or response in valid_responses:
                     break
                 if ord(response) in os_special_keystrokes:
                     break
@@ -48,30 +49,40 @@ 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()
     try:
-        response = _single_keystroke_response_internal(
-            valid_responses, timeout_seconds
-        )
+        response = _single_keystroke_response_internal(valid_responses, timeout_seconds)
         if ord(response) == 3:
             raise KeyboardInterrupt('User pressed ^C in input_utils.')
 
     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:
+def yn_response(prompt: str = None, *, timeout_seconds=None) -> Optional[str]:
     """Get a Y/N response to a prompt."""
 
-    return single_keystroke_response(
+    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..."""
+
+    return single_keystroke_response(None, prompt=prompt, timeout_seconds=timeout_seconds)
 
 
 def keystroke_helper() -> None: