Timed input in Python for Linux systems

timedinput.py

from contextlib import contextmanager
import signal, sys, termios, atexit, tty


class CTLSEQ:
    """
    Condensed from https://github.com/wilfreddv/HB/blob/main/hbutil/hbutil/termctl.py

    Define common ANSI escape
    code control sequences
    """
    ESC = '\x1b'
    SCI = ESC + '['
    BACKSPACE = (chr(0x08), chr(0x7F))
    TERMINATOR = ('\n', '\x04')

    C_LEFT = SCI + "{}D"


class Terminal:
    """
    Condensed from https://github.com/wilfreddv/HB/blob/main/hbutil/hbutil/termctl.py

    Wrapper for some terminal interactions
    """

    _STDIN = sys.stdin.fileno()

    def __init__(self, of=sys.stdout):
        # Save old settings
        self.old_termios = termios.tcgetattr(self._STDIN)
        self.termios = self.old_termios.copy()
        self.of = of
        # Make sure to set the settings back, even after an exception
        atexit.register(self.reset)

    def set_raw(self):
        # Set terminal to "raw" mode
        tty.setcbreak(self._STDIN)

    def read(self, n=1):
        c = self.getch(1)
        if c == CTLSEQ.ESC:
            return c + self.getch(2)
        else:
            # wtf?
            c += self.getch(n-1)
        return c

    def write(self, data):
        self.of.write(data)
        self.of.flush()

    def move_cursor(self, distance, direction):
        self.write(direction.format(distance))

    def getch(self, n=1):
        return sys.stdin.read(n) if n else ""

    def reset(self):
        termios.tcsetattr(self._STDIN,
                          termios.TCSANOW,
                          self.old_termios)


def timedinput(time, prompt=None, default_return=""):
    """
    prompt -> str
    time (seconds) -> int
    """

    @contextmanager
    def timeout():
        def handler(signum, frame):
            raise TimeoutError
        signal.signal(signal.SIGALRM, handler)
        signal.alarm(time)
        try:
            yield
        finally:
            signal.alarm(0)

    terminal = Terminal()
    terminal.set_raw()
    prompt = prompt or ""
    terminal.write(prompt)

    def _input():
        buffer = ""
        try:
            while True:
                char = terminal.read(1)
                if char in CTLSEQ.TERMINATOR:
                    # Done writing, visually seperate and return
                    terminal.write('\n')
                    return buffer
                elif char in CTLSEQ.BACKSPACE:
                    if len(buffer) >= 1:
                        # Remove 1 from buffer, clear previous character
                        buffer = buffer[:-1]
                        terminal.move_cursor(1, CTLSEQ.C_LEFT)
                        terminal.write(" ")
                        terminal.move_cursor(1, CTLSEQ.C_LEFT)
                    else:
                        buffer = ""
                elif char.isprintable():
                    # Only add to the buffer if it makes sense to add
                    buffer += char
                    terminal.write(char)
        except TimeoutError:
            # Time's up!
            terminal.write('\n')
            return buffer or default_return
        finally:
            terminal.reset()

    with timeout():
        return _input()