input_utils.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """Utilities related to user input."""
   6
   7 import logging
   8 import signal
   9 import sys
  10 from typing import List, Optional
  11
  12 import readchar  # type: ignore
  13
  14 import exceptions
  15
  16 logger = logging.getLogger(__file__)
  17
  18
  19 def single_keystroke_response(
  20     valid_responses: Optional[List[str]],  # None = accept anything
  21     *,
  22     prompt: str = None,
  23     default_response: str = None,
  24     timeout_seconds: int = None,
  25 ) -> Optional[str]:  # None if timeout w/o keystroke
  26     """Get a single keystroke response to a prompt."""
  27
  28     def _handle_timeout(signum, frame) -> None:
  29         raise exceptions.TimeoutError()
  30
  31     def _single_keystroke_response_internal(
  32         valid_responses: Optional[List[str]], timeout_seconds: int = None
  33     ) -> str:
  34         os_special_keystrokes = [3, 26]  # ^C, ^Z
  35         if timeout_seconds is not None:
  36             signal.signal(signal.SIGALRM, _handle_timeout)
  37             signal.alarm(timeout_seconds)
  38
  39         try:
  40             while True:
  41                 response = readchar.readchar()
  42                 logger.debug('Keystroke: 0x%x', ord(response))
  43                 if not valid_responses or response in valid_responses:
  44                     break
  45                 if ord(response) in os_special_keystrokes:
  46                     break
  47             return response
  48         finally:
  49             if timeout_seconds is not None:
  50                 signal.alarm(0)
  51
  52     response = None
  53     if prompt is not None:
  54         print(prompt, end="")
  55         sys.stdout.flush()
  56     try:
  57         response = _single_keystroke_response_internal(valid_responses, timeout_seconds)
  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 and response:
  65         print(response)
  66     return response
  67
  68
  69 def yn_response(prompt: str = None, *, timeout_seconds=None) -> Optional[str]:
  70     """Get a Y/N response to a prompt."""
  71
  72     yn = single_keystroke_response(
  73         ["y", "n", "Y", "N"], prompt=prompt, timeout_seconds=timeout_seconds
  74     )
  75     if yn:
  76         yn = yn.lower()
  77     return yn
  78
  79
  80 def press_any_key(
  81     prompt: str = "Press any key to continue...", *, timeout_seconds=None
  82 ) -> Optional[str]:
  83     """Press any key to continue..."""
  84
  85     return single_keystroke_response(None, prompt=prompt, timeout_seconds=timeout_seconds)
  86
  87
  88 def up_down_enter() -> Optional[str]:
  89     os_special_keystrokes = [3, 26]  # ^C, ^Z
  90     while True:
  91         key = readchar.readkey()
  92         if len(key) == 1:
  93             if ord(key) in os_special_keystrokes:
  94                 return None
  95             if ord(key) == 13:
  96                 return 'enter'
  97         elif len(key) == 3:
  98             if ord(key[0]) == 27 and ord(key[1]) == 91:
  99                 if ord(key[2]) == 65:
 100                     return "up"
 101                 elif ord(key[2]) == 66:
 102                     return "down"
 103
 104
 105 def keystroke_helper() -> None:
 106     """Misc util to watch keystrokes and report what they were."""
 107
 108     print("Watching for keystrokes; ^C to quit.")
 109     while True:
 110         key = readchar.readkey()
 111         if len(key) == 1:
 112             print(f'That was "{key}" ({ord(key)}).')
 113             if ord(key) == 3:
 114                 return
 115         else:
 116             print(f'That was sequence "{key}" (', end="")
 117             for _ in key:
 118                 print(f" {ord(_)} ", end="")
 119             print(")")