Python context manager and decorator to timeout after a certain amount of time.

timeout.py

import signal
from contextlib import contextmanager
from typing import Any, Generator


class TimeoutException(Exception):
    """Raised when a timeout occurs."""

    pass


def timeout_handler(signum: int, frame: Any) -> None:
    """Raise a TimeoutException when a signal is received."""
    raise TimeoutException


@contextmanager
def timeout_after_n_sec(sec: int) -> Generator[None, None, None]:
    """Context manager to timeout after sec seconds.

    Args:
        sec (int): number of seconds to timeout after.

    Raises:
        TimeoutException: if the code inside the context manager takes longer
            than sec seconds to run.
    """
    # Set signal handler here to make it isn't overwritten by other
    # signal handlers.
    signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(sec)
    try:
        # Allow the code inside the context manager to run.
        yield
    finally:
        # Cancel the alarm.
        signal.alarm(0)