summaryrefslogtreecommitdiff
path: root/input_utils.py
blob: 0b32eea6f55c5d2fe01466bda0d21ae04c7b13f3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
#!/usr/bin/env python3

# © Copyright 2021-2022, Scott Gasch

"""Utilities related to user input."""

import logging
import signal
import sys
from typing import List, Optional

import readchar  # type: ignore

import exceptions

logger = logging.getLogger(__file__)


def single_keystroke_response(
    valid_responses: Optional[List[str]],  # None = accept anything
    *,
    prompt: str = None,
    default_response: str = None,
    timeout_seconds: int = None,
) -> Optional[str]:  # None if timeout w/o keystroke
    """Get a single keystroke response to a prompt and returns it.

    Args:
        valid_responses: a list of strings that are considered to be
            valid keystrokes to be accepted.  If None, we accept
            anything.
        prompt: the prompt to print before watching keystrokes.  If
            None, skip this.
        default_response: the response to return if the timeout
            expires.  If None, skip this.
        timeout_seconds: number of seconds to wait before timing out
            and returning the default_response.  If None, wait forever.

    Returns:
        The keystroke the user pressed.  If the user pressed a special
        keystroke like ^C or ^Z, we raise a KeyboardInterrupt exception.
    """

    def _handle_timeout(signum, frame) -> None:
        raise exceptions.TimeoutError()

    def _single_keystroke_response_internal(
        valid_responses: Optional[List[str]], timeout_seconds: int = None
    ) -> str:
        os_special_keystrokes = [3, 26]  # ^C, ^Z
        if timeout_seconds is not None:
            signal.signal(signal.SIGALRM, _handle_timeout)
            signal.alarm(timeout_seconds)

        try:
            while True:
                response = readchar.readchar()
                logger.debug('Keystroke: 0x%x', ord(response))
                if not valid_responses or response in valid_responses:
                    break
                if ord(response) in os_special_keystrokes:
                    break
            return response
        finally:
            if timeout_seconds is not None:
                signal.alarm(0)

    response = None
    if prompt is not None:
        print(prompt, end="")
        sys.stdout.flush()
    try:
        response = _single_keystroke_response_internal(valid_responses, timeout_seconds)
        if ord(response) == 3:
            raise KeyboardInterrupt('User pressed ^C in input_utils.')

    except exceptions.TimeoutError:
        if default_response is not None:
            response = default_response
    if prompt and response:
        print(response)
    return response


def yn_response(prompt: str = None, *, timeout_seconds=None) -> Optional[str]:
    """Get a Y/N response to a prompt.

    Args:
        prompt: the user prompt or None to skip this
        timeout_seconds: the number of seconds to wait for a response or
            None to wait forever.

    Returns:
        A lower case 'y' or 'n'.  Or None if the timeout expires with
        no input from the user.  Or raises a KeyboardInterrupt if the
        user pressed a special key such as ^C or ^Z.
    """
    yn = single_keystroke_response(
        ["y", "n", "Y", "N"], prompt=prompt, timeout_seconds=timeout_seconds
    )
    if yn:
        yn = yn.lower()
    return yn


def press_any_key(
    prompt: str = "Press any key to continue...", *, timeout_seconds=None
) -> Optional[str]:
    """Press any key to continue..."""

    return single_keystroke_response(None, prompt=prompt, timeout_seconds=timeout_seconds)


def up_down_enter() -> Optional[str]:
    """Respond to UP, DOWN or ENTER events for simple menus without
    the need for curses."""

    os_special_keystrokes = [3, 26]  # ^C, ^Z
    while True:
        key = readchar.readkey()
        if len(key) == 1:
            if ord(key) in os_special_keystrokes:
                return None
            if ord(key) == 13:
                return 'enter'
        elif len(key) == 3:
            if ord(key[0]) == 27 and ord(key[1]) == 91:
                if ord(key[2]) == 65:
                    return "up"
                elif ord(key[2]) == 66:
                    return "down"


def keystroke_helper() -> None:
    """Misc util to watch keystrokes and report what they were."""

    print("Watching for keystrokes; ^C to quit.")
    while True:
        key = readchar.readkey()
        if len(key) == 1:
            print(f'That was "{key}" ({ord(key)}).')
            if ord(key) == 3:
                return
        else:
            print(f'That was sequence "{key}" (', end="")
            for _ in key:
                print(f" {ord(_)} ", end="")
            print(")")