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