Thread Pool Map With Timeout

import time
from collections.abc import Callable, Sequence
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout
from concurrent.futures import as_completed
from typing import TypeVar

T = TypeVar("T")
R = TypeVar("R")


def pool_map(
    func: Callable[[T], R],
    items: Sequence[T],
    *,
    workers: int = 4,
    timeout: float | None = None,
) -> list[R]:
    if workers < 1:
        raise ValueError("workers must be positive")
    if not items:
        return []
    results: dict[int, R] = {}
    with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="map") as pool:
        futures = {pool.submit(func, item): index for index, item in enumerate(items)}
        try:
            for future in as_completed(futures, timeout=timeout):
                results[futures[future]] = future.result()
        except FutureTimeout:
            for future in futures:
                future.cancel()
            raise
    return [results[index] for index in range(len(items))]


def main() -> None:
    def slow_square(value: int) -> int:
        time.sleep(0.05)
        return value * value

    numbers = list(range(1, 7))
    start = time.perf_counter()
    squares = pool_map(slow_square, numbers, workers=3)
    print("squares:", squares)
    print(f"elapsed: {time.perf_counter() - start:.2f}s")

    try:
        pool_map(lambda value: time.sleep(0.3) or value, [1, 2], timeout=0.05)
    except FutureTimeout:
        print("timed out and cancelled pending work")


if __name__ == "__main__":
    main()