summaryrefslogtreecommitdiff
path: root/input_utils.py
blob: 7d5e18089331a96bc32172b90e916a37c0695d49 (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
#!/usr/bin/env python3

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

import logging
import signal
import sys
from typing import List

import readchar  # type: ignore

import exceptions

logger = logging.getLogger(__file__)


def single_keystroke_response(
    valid_responses: List[str],
    *,
    prompt: str = None,
    default_response: str = None,
    timeout_seconds: int = None,
) -> str:
    """Get a single keystroke response to a prompt."""

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

    def _single_keystroke_response_internal(
        valid_responses: List[str], timeout_seconds=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 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)

    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 is not None:
        print(response)
    return response


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

    return single_keystroke_response(
        ["y", "n", "Y", "N"], prompt=prompt, timeout_seconds=timeout_seconds
    ).lower()


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(")")