timeout.py

download raw

import threading
import _thread
import typing

T = typing.TypeVar("T")
P = typing.ParamSpec("P")


def run_with_timeout(
    timeout: float, function: typing.Callable[P, T], /, *args: P.args, **kw: P.kwargs
) -> tuple[typing.Literal[False], KeyboardInterrupt] | tuple[typing.Literal[True], T]:
    """Run a function with a timeout.

    This has various assumptions including:
     * the function does not capture KeyboardInterupt,
     * the kind of activity it is performing can be interrupted by `_thread.interrupt_main`
     * that `run_with_timeout` is called from the main thread."""

    timer = threading.Timer(timeout, _thread.interrupt_main)
    try:
        try:
            timer.start()
            return True, function(*args, **kw)
        finally:
            timer.cancel()
    except KeyboardInterrupt as e:
        return False, e


if __name__ == "__main__":
    import time

    def fib_recursive_bad(n):
        if n < 2:
            return n
        return fib_recursive_bad(n - 1) + fib_recursive_bad(n - 2)

    for n in range(20, 800):
        t0 = time.monotonic()
        ok, result = run_with_timeout(1.25, fib_recursive_bad, n)
        t1 = time.monotonic()
        dt = t1 - t0
        print(f"fib({n}) = {result!r} (in {dt*1000:.0f} ms)")
        if not ok:
            break