input_utils.py

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