How to get next itereration of async generator after calling `async for in`

user_532963 0
Viewed 1226

Using FastAPI I am trying to detect if a StreamingResponse has been entirely been consumed by the client or if it was cancelled.

I have the following example app:

import asyncio

import uvicorn
from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()


async def ainfinite_generator():
    while True:
        yield b"some fake data "
        await asyncio.sleep(.001)


async def astreamer(generator):
    try:
        async for data in generator:
            yield data
    except Exception as e:
        # this isn't triggered by a cancelled request
        print(e)
    finally:
        # this always throws a StopAsyncIteration exception
        # no matter whether the generator was consumed or not
        leftover = await generator.__anext__()
        if leftover:
            print("we didn't finish")
        else:
            print("we finished")


@app.get("/")
async def infinite_stream():
    return StreamingResponse(astreamer(ainfinite_generator()))


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

It seems like the first async for in generator in astreamer "consumes" the async generator. After that loop, further attempts to get the next iteration fail with a StopAsyncIteration exception, even if the generator is "infinite" as defined above.

I've looked through PEP-525 and the only thing I am seeing is that if an exception is thrown into the generator it will cause any further attempts to read from the generator to throw that StopAsyncIteration exception, but I don't see where that would be happening. At least, I'm not seeing that in Starlette's StreamingResponse class (it doesn't seem to do much with "content"). Does the generator not get "released" after doing an async for in gen?

starlette fastapi python-asyncio python
1 Answers

The code below shows how to watch for cancellation on coroutines (in my case, an async generators). As mentioned in the comments, if an async generator is cancelled, it injects an exception into the generator, and from that point on, any attempt to get the next item in the generator will raise the StopAsyncIteration exception. See PEP 525. To determine if the async generator was cancelled, just try/except on a asyncio.CancelledError exception (which derives from BaseException).

There's also code here for showing how to deal with normal generators, which are a bit more forgiving. If you keep the same try/except flow, those raise a GeneratorExit exception if they are cancelled.

The tricky part was that most of these exceptions derive from the BaseException class, unlike the StopIteration exception I was expecting, which derives from Exception class.

And, by the way, the actual cancellation is happening in starlette.

import asyncio
import time

import uvicorn
from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()


def infinite_generator():
    # not blocking, so doesn't need to be async
    # but if it was blocking, you could make this async and await it
    while True:
        yield b"some fake data "


def finite_generator():
    # not blocking, so doesn't need to be async
    # but if it was blocking, you could make this async and await it
    x = 0
    while x < 10000:
        yield f"{x}"
        x += 1


async def astreamer(generator):
    try:
        # if it was an async generator we'd do:
        # "async for data in generator:"
        # (there is no yield from async_generator)
        for i in generator:
            yield i
            await asyncio.sleep(.001)

    except asyncio.CancelledError as e:
        print('cancelled')


def streamer(generator):
    try:
        # note: normally we would do "yield from generator"
        # but that won't work with next(generator) in the finally statement
        for i in generator:
            yield i
            time.sleep(.001)

    except GeneratorExit:
        print("cancelled")
    finally:
        # showing that we can check here to see if all data was consumed
        # the except statement above effectively does the same thing
        try:
            next(generator)
            print("we didn't finish")
            return
        except StopIteration:
            print("we finished")


@app.get("/infinite")
async def infinite_stream():
    return StreamingResponse(streamer(infinite_generator()))


@app.get("/finite")
async def finite_stream():
    return StreamingResponse(streamer(finite_generator()))


@app.get("/ainfinite")
async def infinite_stream():
    return StreamingResponse(astreamer(infinite_generator()))


@app.get("/afinite")
async def finite_stream():
    return StreamingResponse(astreamer(finite_generator()))


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
user_5329630
Related