from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor


def timeout(*args, **kwargs):
    """
    Limit the function execution time to the given timeout (in seconds).

    Example:
        >>> fib = lambda n: 1 if n <= 1 else fib(n - 1) + fib(n - 2)
        >>> timeout(0.25, fib, 10)  # wait at most 0.25 seconds
        89
        >>> timeout(0.25, fib, 50)  # stops before the thermal death of the universe
        Traceback (most recent call last)
        ...
        TimeoutError:

    Args:
        timeout (float, seconds):
            The maximum duration of function execution
        func:
            Function to be executed
    """

    executor = ProcessPoolExecutor
    timeout, func, *args = args
    if func in ("thread", "process"):
        executor = ThreadPoolExecutor if func == "thread" else ProcessPoolExecutor
        func, *args = args

    with executor() as e:
        future = e.submit(func, *args, **kwargs)
        return future.result(timeout=timeout)