input_utils.py

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