1"""
2A snippet for running multiple, concurrent invocations of a Python function.
3"""
4
5from collections.abc import Callable, Iterable
6import concurrent.futures
7import itertools
8import typing
9
10
11In = typing.TypeVar("In")
12Out = typing.TypeVar("Out")
13
14
15def concurrently(
16    handler: Callable[[In], Out], inputs: Iterable[In], *, max_concurrency: int = 5
17) -> Iterable[tuple[In, Out]]:
18    """
19    Call the function ``handler`` on the values ``inputs``.
20
21    ``handler`` should be a function that takes a single input, which is the
22    individual values in the iterable ``inputs``.
23
24    Generates (input, output) tuples as the calls to ``handler`` complete.
25
26    See https://alexwlchan.net/2019/10/adventures-with-concurrent-futures/
27    for an explanation of how this function works.
28
29    """
30    # Make sure we get a consistent iterator throughout, rather than
31    # getting the first element repeatedly.
32    handler_inputs = iter(inputs)
33
34    with concurrent.futures.ThreadPoolExecutor() as executor:
35        futures = {
36            executor.submit(handler, input): input
37            for input in itertools.islice(handler_inputs, max_concurrency)
38        }
39
40        while futures:
41            done, _ = concurrent.futures.wait(
42                futures, return_when=concurrent.futures.FIRST_COMPLETED
43            )
44
45            for fut in done:
46                original_input = futures.pop(fut)
47                yield original_input, fut.result()
48
49            for input in itertools.islice(handler_inputs, len(done)):
50                fut = executor.submit(handler, input)
51                futures[fut] = input