input_utils.py

   1 #!/usr/bin/env python3
   2
   3 """Utilities related to user input."""
   4
   5 import signal
   6 import sys
   7 from typing import List
   8
   9 import readchar  # type: ignore
  10
  11 import exceptions
  12
  13
  14 def single_keystroke_response(
  15     valid_responses: List[str],
  16     *,
  17     prompt: str = None,
  18     default_response: str = None,
  19     timeout_seconds: int = None,
  20 ) -> str:
  21     """Get a single keystroke response to a prompt."""
  22
  23     def _handle_timeout(signum, frame) -> None:
  24         raise exceptions.TimeoutError()
  25
  26     def _single_keystroke_response_internal(
  27         valid_responses: List[str], timeout_seconds=None
  28     ) -> str:
  29         os_special_keystrokes = [3, 26]  # ^C, ^Z
  30         if timeout_seconds is not None:
  31             signal.signal(signal.SIGALRM, _handle_timeout)
  32             signal.alarm(timeout_seconds)
  33
  34         try:
  35             while True:
  36                 response = readchar.readchar()
  37                 if response in valid_responses:
  38                     break
  39                 if ord(response) in os_special_keystrokes:
  40                     break
  41             return response
  42         finally:
  43             if timeout_seconds is not None:
  44                 signal.alarm(0)
  45
  46     if prompt is not None:
  47         print(prompt, end="")
  48         sys.stdout.flush()
  49     try:
  50         response = _single_keystroke_response_internal(
  51             valid_responses, timeout_seconds
  52         )
  53     except exceptions.TimeoutError:
  54         if default_response is not None:
  55             response = default_response
  56     if prompt is not None:
  57         print(response)
  58     return response
  59
  60
  61 def yn_response(prompt: str = None, *, timeout_seconds=None) -> str:
  62     """Get a Y/N response to a prompt."""
  63
  64     return single_keystroke_response(
  65         ["y", "n", "Y", "N"], prompt=prompt, timeout_seconds=timeout_seconds
  66     ).lower()
  67
  68
  69 def keystroke_helper() -> None:
  70     """Misc util to watch keystrokes and report what they were."""
  71
  72     print("Watching for keystrokes; ^C to quit.")
  73     while True:
  74         key = readchar.readkey()
  75         if len(key) == 1:
  76             print(f'That was "{key}" ({ord(key)}).')
  77             if ord(key) == 3:
  78                 return
  79         else:
  80             print(f'That was sequence "{key}" (', end="")
  81             for _ in key:
  82                 print(f" {ord(_)} ", end="")
  83             print(")")