asyncio: coroutines and tasks

import asyncio
from collections.abc import Awaitable, Callable


async def fetch_all(names: list[str], fetch: Callable[[str], Awaitable[str]]) -> dict[str, str]:
    """Run fetch(name) for every name concurrently; map each name to its text, in the order of names."""
    # Each await waits for one fetch before the next one starts.
    # Start them all together with asyncio.gather instead.
    return {name: await fetch(name) for name in names}

import asyncio

from main import fetch_all


def test_results():
    """Every name maps to its text"""

    async def fake_fetch(name):
        await asyncio.sleep(0)
        return name * 2

    got = asyncio.run(fetch_all(["a", "b"], fake_fetch))
    assert got == {"a": "aa", "b": "bb"}, f"fetch_all returned {got!r}"


def test_order_of_names():
    """The dict follows the order of names, even when a later fetch finishes first"""

    async def fake_fetch(name):
        await asyncio.sleep(0.02 if name == "slow" else 0)
        return name.upper()

    got = asyncio.run(fetch_all(["slow", "fast"], fake_fetch))
    assert list(got.items()) == [("slow", "SLOW"), ("fast", "FAST")], f"fetch_all returned {got!r}"


def test_fetches_overlap():
    """All three fetches are waiting at the same time"""
    running = 0
    peak = 0

    async def fake_fetch(name):
        nonlocal running, peak
        running += 1
        peak = max(peak, running)
        await asyncio.sleep(0)
        running -= 1
        return name

    asyncio.run(fetch_all(["a", "b", "c"], fake_fetch))
    assert peak == 3, f"at most {peak} fetch(es) waited at once: start them all with asyncio.gather"

import asyncio
import time


async def brew(name: str, seconds: float, log: list[str]) -> str:
    """Pretend to brew a drink for some seconds; log when it starts and ends."""
    log.append(f"{name} start")
    time.sleep(seconds)  # blocks the event loop: no other task can run meanwhile
    log.append(f"{name} end")
    return f"{name} ready"


async def brew_all(orders: dict[str, float]) -> tuple[list[str], list[str]]:
    """Brew every order in its own task; return (results in order, log)."""
    log: list[str] = []
    tasks = [asyncio.create_task(brew(name, seconds, log)) for name, seconds in orders.items()]
    results = [await task for task in tasks]
    return results, log

import asyncio

from main import brew_all


def test_results():
    """Every order is ready, in the order of the orders"""
    results, _ = asyncio.run(brew_all({"tea": 0.01, "coffee": 0}))
    assert results == ["tea ready", "coffee ready"], f"brew_all returned {results!r}"


def test_brews_overlap():
    """Coffee starts while tea is still brewing, and the shorter brew ends first"""
    _, log = asyncio.run(brew_all({"tea": 0.02, "coffee": 0.01}))
    assert log == ["tea start", "coffee start", "coffee end", "tea end"], f"the log was {log!r}: does brew block the event loop?"

import asyncio
import sys


async def run_python(code: str) -> tuple[int, str]:
    """Run code in a new Python process; return its exit status and its output without the last newline."""
    # TODO: start sys.executable with the arguments "-c" and code, standard output piped;
    # await communicate(), then return the exit status and the decoded output.
    raise NotImplementedError


async def run_all(snippets: list[str]) -> list[tuple[int, str]]:
    """Run every snippet at the same time; return the results in the order of snippets."""
    # This waits for each program before it starts the next one.
    return [await run_python(code) for code in snippets]


if __name__ == "__main__":
    print(asyncio.run(run_all(["print('hello')", "import sys; sys.exit(3)"])))

import asyncio
import time

from main import run_all, run_python


def test_output():
    """run_python returns exit status 0 and what the program printed"""
    got = asyncio.run(run_python("print('hi')"))
    assert got == (0, "hi"), f"run_python returned {got!r}, expected (0, 'hi')"


def test_exit_status():
    """A program that ends with sys.exit(3) gives exit status 3"""
    got = asyncio.run(run_python("import sys; sys.exit(3)"))
    assert got == (3, ""), f"run_python returned {got!r}, expected (3, '')"


def test_programs_overlap():
    """Three programs of 0.5 s each run at the same time, and the results keep their order"""
    snippets = [f"import time; time.sleep(0.5); print({n})" for n in range(3)]
    started = time.perf_counter()
    got = asyncio.run(run_all(snippets))
    took = time.perf_counter() - started
    assert got == [(0, "0"), (0, "1"), (0, "2")], f"run_all returned {got!r}"
    assert took < 1.4, f"run_all took {took:.2f} s: start every program before you wait for any"