TaskGroup, timeouts and cancellation

import asyncio
from collections.abc import Callable, Coroutine
from typing import Any


async def fetch_all(names: list[str], fetch: Callable[[str], Coroutine[Any, Any, str]]) -> dict[str, str]:
    """Fetch every name in a TaskGroup; map each name to its text, in the order of names.

    If a fetch fails, the others are cancelled and an ExceptionGroup is raised.
    """
    # gather leaves the other fetches running when one fails.
    # Create one task per name in an asyncio.TaskGroup instead.
    texts = await asyncio.gather(*(fetch(name) for name in names))
    return dict(zip(names, texts))

import asyncio

from main import fetch_all


def test_results():
    """Every name maps to its text, in the order of names"""

    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_failure_is_a_group():
    """A failing fetch surfaces as an ExceptionGroup holding its ValueError"""

    async def fake_fetch(name):
        await asyncio.sleep(0)
        if name == "bad":
            raise ValueError("bad page")
        return name

    try:
        asyncio.run(fetch_all(["ok", "bad"], fake_fetch))
        error = None
    except Exception as caught:
        error = caught
    assert isinstance(error, ExceptionGroup), f"fetch_all raised {error!r}, not an ExceptionGroup"
    assert [str(e) for e in error.exceptions] == ["bad page"], f"the group holds {error.exceptions!r}"


def test_others_cancelled():
    """When one fetch fails, the slow one is cancelled before fetch_all raises"""
    cancelled = []

    async def fake_fetch(name):
        if name == "bad":
            await asyncio.sleep(0)
            raise ValueError("bad page")
        try:
            await asyncio.sleep(1)
        except asyncio.CancelledError:
            cancelled.append(name)
            raise
        return name

    async def run():
        try:
            await fetch_all(["slow", "bad"], fake_fetch)
        except Exception:
            pass
        return list(cancelled)

    seen = asyncio.run(run())
    assert seen == ["slow"], f"cancelled when fetch_all raised: {seen!r}; use a TaskGroup"

import asyncio
from collections.abc import Awaitable, Callable


async def fetch_one(fetch: Callable[[str], Awaitable[str]], name: str, limit: float) -> str:
    """Return fetch(name), or "timeout" if it takes longer than limit seconds."""
    async with asyncio.timeout(limit):
        try:
            return await fetch(name)
        except TimeoutError:  # never true here: inside the block the fetch sees CancelledError
            return "timeout"


async def fetch_each(names: list[str], fetch: Callable[[str], Awaitable[str]], limit: float) -> list[str]:
    """Fetch every name in a TaskGroup, each with its own time limit; results in the order of names."""
    async with asyncio.TaskGroup() as tg:
        tasks = [tg.create_task(fetch_one(fetch, name, limit)) for name in names]
    return [task.result() for task in tasks]

import asyncio

from main import fetch_each

DELAYS = {"home": 0, "blog": 0.01, "archive": 1}


async def fake_fetch(name):
    await asyncio.sleep(DELAYS[name])
    return name.upper()


def test_all_in_time():
    """Pages that arrive in time give their text"""
    got = asyncio.run(fetch_each(["home", "blog"], fake_fetch, 0.5))
    assert got == ["HOME", "BLOG"], f"fetch_each returned {got!r}"


def test_slow_page_times_out():
    """A page slower than the limit gives "timeout", and the others still arrive"""
    try:
        got = asyncio.run(fetch_each(["home", "archive", "blog"], fake_fetch, 0.05))
    except Exception as error:
        got = error
    assert got == ["HOME", "timeout", "BLOG"], f"fetch_each gave {got!r}: where is the TimeoutError raised?"

import asyncio
from collections.abc import Awaitable, Callable


async def fetch_limited(names: list[str], fetch: Callable[[str], Awaitable[str]], limit: int) -> list[str]:
    """Fetch every name in a TaskGroup, at most limit at a time; results in the order of names."""
    # TODO: every fetch starts at once. Create one asyncio.Semaphore(limit)
    # here and hold it in fetch_one around the await.

    async def fetch_one(name: str) -> str:
        return await fetch(name)

    async with asyncio.TaskGroup() as tg:
        tasks = [tg.create_task(fetch_one(name)) for name in names]
    return [task.result() for task in tasks]


async def demo_fetch(name: str) -> str:
    print("start", name)
    await asyncio.sleep(0.1)
    print("end", name)
    return name.upper()


if __name__ == "__main__":
    print(asyncio.run(fetch_limited(["home", "about", "blog", "shop"], demo_fetch, 2)))

import asyncio

from main import fetch_limited


def counting_fetch():
    """A fake fetch that records how many calls are waiting at the same time."""
    state = {"running": 0, "peak": 0}

    async def fetch(name):
        state["running"] += 1
        state["peak"] = max(state["peak"], state["running"])
        await asyncio.sleep(0.01)
        state["running"] -= 1
        return name.upper()

    return fetch, state


def test_results_in_order():
    """Every name gives its text, in the order of names"""
    fetch, _ = counting_fetch()
    got = asyncio.run(fetch_limited(["home", "about", "blog"], fetch, 2))
    assert got == ["HOME", "ABOUT", "BLOG"], f"fetch_limited returned {got!r}"


def test_limit_is_kept():
    """With limit 2, never more than 2 fetches run at once"""
    fetch, state = counting_fetch()
    asyncio.run(fetch_limited(["a", "b", "c", "d", "e"], fetch, 2))
    assert state["peak"] == 2, f"{state['peak']} fetches ran at once, expected at most 2: hold one shared Semaphore around each fetch"


def test_limit_is_used():
    """With limit 3, three fetches do run at the same time"""
    fetch, state = counting_fetch()
    asyncio.run(fetch_limited(["a", "b", "c", "d", "e"], fetch, 3))
    assert state["peak"] == 3, f"{state['peak']} fetches ran at once, expected 3: create the Semaphore once, not in every fetch"