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 and returns it.
  27
  28     Args:
  29         valid_responses: a list of strings that are considered to be
  30             valid keystrokes to be accepted.  If None, we accept
  31             anything.
  32         prompt: the prompt to print before watching keystrokes.  If
  33             None, skip this.
  34         default_response: the response to return if the timeout
  35             expires.  If None, skip this.
  36         timeout_seconds: number of seconds to wait before timing out
  37             and returning the default_response.  If None, wait forever.
  38
  39     Returns:
  40         The keystroke the user pressed.  If the user pressed a special
  41         keystroke like ^C or ^Z, we raise a KeyboardInterrupt exception.
  42     """
  43
  44     def _handle_timeout(signum, frame) -> None:
  45         raise exceptions.TimeoutError()
  46
  47     def _single_keystroke_response_internal(
  48         valid_responses: Optional[List[str]], timeout_seconds: int = None
  49     ) -> str:
  50         os_special_keystrokes = [3, 26]  # ^C, ^Z
  51         if timeout_seconds is not None:
  52             signal.signal(signal.SIGALRM, _handle_timeout)
  53             signal.alarm(timeout_seconds)
  54
  55         try:
  56             while True:
  57                 response = readchar.readchar()
  58                 logger.debug('Keystroke: 0x%x', ord(response))
  59                 if not valid_responses or response in valid_responses:
  60                     break
  61                 if ord(response) in os_special_keystrokes:
  62                     break
  63             return response
  64         finally:
  65             if timeout_seconds is not None:
  66                 signal.alarm(0)
  67
  68     response = None
  69     if prompt is not None:
  70         print(prompt, end="")
  71         sys.stdout.flush()
  72     try:
  73         response = _single_keystroke_response_internal(valid_responses, timeout_seconds)
  74         if ord(response) == 3:
  75             raise KeyboardInterrupt('User pressed ^C in input_utils.')
  76
  77     except exceptions.TimeoutError:
  78         if default_response is not None:
  79             response = default_response
  80     if prompt and response:
  81         print(response)
  82     return response
  83
  84
  85 def yn_response(prompt: str = None, *, timeout_seconds=None) -> Optional[str]:
  86     """Get a Y/N response to a prompt.
  87
  88     Args:
  89         prompt: the user prompt or None to skip this
  90         timeout_seconds: the number of seconds to wait for a response or
  91             None to wait forever.
  92
  93     Returns:
  94         A lower case 'y' or 'n'.  Or None if the timeout expires with
  95         no input from the user.  Or raises a KeyboardInterrupt if the
  96         user pressed a special key such as ^C or ^Z.
  97     """
  98     yn = single_keystroke_response(
  99         ["y", "n", "Y", "N"], prompt=prompt, timeout_seconds=timeout_seconds
 100     )
 101     if yn:
 102         yn = yn.lower()
 103     return yn
 104
 105
 106 def press_any_key(
 107     prompt: str = "Press any key to continue...", *, timeout_seconds=None
 108 ) -> Optional[str]:
 109     """Press any key to continue..."""
 110
 111     return single_keystroke_response(None, prompt=prompt, timeout_seconds=timeout_seconds)
 112
 113
 114 def up_down_enter() -> Optional[str]:
 115     """Respond to UP, DOWN or ENTER events for simple menus without
 116     the need for curses."""
 117
 118     os_special_keystrokes = [3, 26]  # ^C, ^Z
 119     while True:
 120         key = readchar.readkey()
 121         if len(key) == 1:
 122             if ord(key) in os_special_keystrokes:
 123                 return None
 124             if ord(key) == 13:
 125                 return 'enter'
 126         elif len(key) == 3:
 127             if ord(key[0]) == 27 and ord(key[1]) == 91:
 128                 if ord(key[2]) == 65:
 129                     return "up"
 130                 elif ord(key[2]) == 66:
 131                     return "down"
 132
 133
 134 def keystroke_helper() -> None:
 135     """Misc util to watch keystrokes and report what they were."""
 136
 137     print("Watching for keystrokes; ^C to quit.")
 138     while True:
 139         key = readchar.readkey()
 140         if len(key) == 1:
 141             print(f'That was "{key}" ({ord(key)}).')
 142             if ord(key) == 3:
 143                 return
 144         else:
 145             print(f'That was sequence "{key}" (', end="")
 146             for _ in key:
 147                 print(f" {ord(_)} ", end="")
 148             print(")")