bpo-29842: Make Executor.map less eager so it handles large/unbounded… by graingert · Pull Request #18566 · python/cpython

import contextlib
import functools
import concurrent.futures
import threading
import sys


def fn(num, stop_event):
    if num == 1:
        stop_event.wait()
        return "done 1"

    if num == 2:
        return "done 2"


def main():
    stop_event = threading.Event()
    log = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:

        def print_n_wait(ident):
            log.append(f"{ident=} started")
            try:
                stop_event.wait()
            finally:
                log.append(f"{ident=} stopped")

        fut = pool.submit(print_n_wait, ident="first")
        try:
            with contextlib.closing(pool.map(print_n_wait, ["second", "third"], timeout=1)) as gen:
                try:
                    next(gen)
                except concurrent.futures.TimeoutError:
                    print("timed out")
                else:
                    raise RuntimeError("timeout expected")
        finally:
            stop_event.set()

    assert log == ["ident='first' started", "ident='first' stopped"], f"{log=} is wrong"

if __name__ == "__main__":
    sys.exit(main())