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