Cancellation and Exception Groups

import anyio


async def stream_answer(log: list) -> None:
    try:
        for token in ["Struct", "ured", " concurrency"]:
            await anyio.sleep(0.02)
            log.append(token)
    except anyio.get_cancelled_exc_class():
        log.append("<client disconnected>")
        raise                                   # always


async def main() -> None:
    log: list = []
    with anyio.move_on_after(0.03):
        await stream_answer(log)
    print(log)


anyio.run(main)
# Output: ['Struct', '<client disconnected>']

import anyio


class VectorBuffer:
    def __init__(self) -> None:
        self.pending = [1, 2, 3]

    async def flush(self) -> None:
        await anyio.sleep(0.01)
        print("flushed", len(self.pending))
        self.pending.clear()


async def ingest(buf: VectorBuffer, shielded: bool) -> None:
    try:
        await anyio.sleep(10)
    finally:
        if shielded:
            with anyio.CancelScope(shield=True), anyio.move_on_after(2):
                await buf.flush()
        else:
            await buf.flush()                    # raises the cancel exception right away


async def main() -> None:
    for shielded in (False, True):
        buf = VectorBuffer()
        with anyio.move_on_after(0.01):
            await ingest(buf, shielded)
        print("shielded" if shielded else "unshielded", "pending:", len(buf.pending))


anyio.run(main)
# Output:
# unshielded pending: 3
# flushed 3
# shielded pending: 0

import anyio


async def main() -> None:
    # 1. outer fail_after fires while the inner move_on_after is still pending
    try:
        with anyio.fail_after(0.05):
            with anyio.move_on_after(1) as inner:
                await anyio.sleep(1)
    except TimeoutError:
        print("outer fired; inner.cancelled_caught =", inner.cancelled_caught)

    # 2. an outer move_on_after cancels; the inner fail_after does NOT raise TimeoutError
    with anyio.move_on_after(0.02) as outer:
        try:
            with anyio.fail_after(5):
                await anyio.sleep(1)
        except TimeoutError:
            print("not printed: the inner deadline did not pass")
    print("outer.cancelled_caught =", outer.cancelled_caught)


anyio.run(main)
# Output:
# outer fired; inner.cancelled_caught = False
# outer.cancelled_caught = True

import anyio


async def call_tool(name: str, fail: bool, log: list) -> None:
    try:
        await anyio.sleep(0.01 if fail else 1)
        if fail:
            raise ConnectionError(f"{name}: upstream reset")
    except anyio.get_cancelled_exc_class():
        log.append(f"{name} cancelled")
        raise


async def main() -> None:
    log: list = []
    try:
        async with anyio.create_task_group() as tg:
            tg.start_soon(call_tool, "web_search", True, log)
            tg.start_soon(call_tool, "sql", False, log)
            tg.start_soon(call_tool, "calculator", False, log)
    except* ConnectionError as eg:
        print("failed:", [str(e) for e in eg.exceptions])
    print(sorted(log))


anyio.run(main)
# Output:
# failed: ['web_search: upstream reset']
# ['calculator cancelled', 'sql cancelled']

import anyio


async def bad(msg: str) -> None:
    await anyio.sleep(0)
    raise ValueError(msg)


async def main() -> None:
    try:
        async with anyio.create_task_group() as tg:
            tg.start_soon(bad, "outer child")
            async with anyio.create_task_group() as inner:
                inner.start_soon(bad, "inner child")
    except* ValueError as eg:
        leaves = []

        def walk(group) -> None:
            for exc in group.exceptions:
                walk(exc) if isinstance(exc, BaseExceptionGroup) else leaves.append(str(exc))

        walk(eg)
        print(sorted(leaves))


anyio.run(main)
# Output: ['inner child', 'outer child']

# Python 3.10 (and newer). Not run here: the backport is not installed in this environment.
from exceptiongroup import ExceptionGroup, catch


def on_rate_limit(eg: ExceptionGroup) -> None:
    print("rate limited", len(eg.exceptions), "calls")


with catch({RateLimitError: on_rate_limit}):
    anyio.run(fan_out, prompts)

import anyio


def unwrap_single(eg: BaseExceptionGroup) -> BaseException:
    """Return the only leaf exception of a (possibly nested) group, or the group itself."""
    leaves = []
    stack = [eg]
    while stack:
        g = stack.pop()
        for exc in g.exceptions:
            (stack if isinstance(exc, BaseExceptionGroup) else leaves).append(exc)
    return leaves[0] if len(leaves) == 1 else eg


class ProviderError(Exception):
    pass


async def call(prompt: str) -> None:
    await anyio.sleep(0)
    raise ProviderError(f"429 for {prompt!r}")


async def public_api(prompts: list) -> None:
    try:
        async with anyio.create_task_group() as tg:
            for p in prompts:
                tg.start_soon(call, p)
    except BaseExceptionGroup as eg:
        exc = unwrap_single(eg)
        if exc is eg:
            raise
        raise exc from eg                       # keep the group as the cause for debugging


async def main() -> None:
    try:
        await public_api(["only one"])
    except ProviderError as exc:
        print("caller sees:", exc)


anyio.run(main)
# Output: caller sees: 429 for 'only one'

async def tokens(prompt):
    async with anyio.create_task_group() as tg:      # scope spans a yield
        tg.start_soon(keepalive)
        async for t in llm.stream(prompt):
            yield t

import anyio
from anyio.streams.memory import MemoryObjectSendStream


async def produce_tokens(prompt: str, send: MemoryObjectSendStream) -> None:
    async with send:
        for token in prompt.split():
            await anyio.sleep(0.001)          # stands in for the model stream
            await send.send(token)


async def main() -> None:
    send, receive = anyio.create_memory_object_stream[str](8)
    async with anyio.create_task_group() as tg:
        tg.start_soon(produce_tokens, "tokens flow through a stream", send)
        async with receive:
            print([t async for t in receive])


anyio.run(main)
# Output: ['tokens', 'flow', 'through', 'a', 'stream']

import time

import anyio


def chunk_corpus(n: int) -> int:
    done = 0
    for _ in range(n):
        time.sleep(0.01)                          # blocking work
        anyio.from_thread.check_cancelled()       # cooperative cancellation point
        done += 1
    return done


async def main() -> None:
    start = time.monotonic()
    with anyio.move_on_after(0.05) as scope:
        await anyio.to_thread.run_sync(chunk_corpus, 1000)
    print(scope.cancelled_caught, time.monotonic() - start < 0.5)


anyio.run(main)
# Output: True True