如何在调用 `async for in` 后获取异步生成器的下一次迭代

Question

使用 FastAPI 我试图检测StreamingResponse是否已完全被客户端使用或是否已被取消。

我有以下示例应用程序：

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)

它似乎是第一个async for in generator“astreamer消耗”异步生成器的。在该循环之后，进一步尝试获得下一次迭代失败并出现异常StopAsyncIteration，即使生成器如上定义的那样是“无限的”。

我查看了PEP-525，我唯一看到的是，如果将异常抛入生成器，它将导致进一步尝试从生成器读取以抛出 StopAsyncIteration 异常，但我看不到它在哪里会发生。至少，我没有在 Starlette 的 StreamingResponse类中看到这一点（它似乎与“内容”无关）。执行后生成器不会“释放”吗async for in gen？

Answer 1

下面的代码展示了如何在协程上观察取消（在我的例子中，是一个异步生成器）。如评论中所述，如果取消异步生成器，它会向生成器注入异常，从那时起，任何试图获取生成器中的下一项的尝试都会引发异常StopAsyncIteration。请参阅PEP 525。要确定异步生成器是否被取消，只需尝试/排除异常asyncio.CancelledError（派生自BaseException）。这里还有一些代码用于展示如何处理普通生成器，这些生成器更宽容一些。如果您保持相同的 try/except 流程，那么在它们GeneratorExit被取消时会引发异常。棘手的部分是这些异常中的大多数都派生自类BaseException，这与StopIteration我期望的派生Exception自类的异常不同。而且，顺便说一句，实际取消发生在starlette中。import asyncioimport timeimport uvicornfrom fastapi import FastAPIfrom fastapi.responses import StreamingResponseapp = FastAPI()def infinite_generator():&nbsp; &nbsp; # not blocking, so doesn't need to be async&nbsp; &nbsp; # but if it was blocking, you could make this async and await it&nbsp; &nbsp; while True:&nbsp; &nbsp; &nbsp; &nbsp; yield b"some fake data "def finite_generator():&nbsp; &nbsp; # not blocking, so doesn't need to be async&nbsp; &nbsp; # but if it was blocking, you could make this async and await it&nbsp; &nbsp; x = 0&nbsp; &nbsp; while x < 10000:&nbsp; &nbsp; &nbsp; &nbsp; yield f"{x}"&nbsp; &nbsp; &nbsp; &nbsp; x += 1async def astreamer(generator):&nbsp; &nbsp; try:&nbsp; &nbsp; &nbsp; &nbsp; # if it was an async generator we'd do:&nbsp; &nbsp; &nbsp; &nbsp; # "async for data in generator:"&nbsp; &nbsp; &nbsp; &nbsp; # (there is no yield from async_generator)&nbsp; &nbsp; &nbsp; &nbsp; for i in generator:&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; yield i&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; await asyncio.sleep(.001)&nbsp; &nbsp; except asyncio.CancelledError as e:&nbsp; &nbsp; &nbsp; &nbsp; print('cancelled')def streamer(generator):&nbsp; &nbsp; try:&nbsp; &nbsp; &nbsp; &nbsp; # note: normally we would do "yield from generator"&nbsp; &nbsp; &nbsp; &nbsp; # but that won't work with next(generator) in the finally statement&nbsp; &nbsp; &nbsp; &nbsp; for i in generator:&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; yield i&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; time.sleep(.001)&nbsp; &nbsp; except GeneratorExit:&nbsp; &nbsp; &nbsp; &nbsp; print("cancelled")&nbsp; &nbsp; finally:&nbsp; &nbsp; &nbsp; &nbsp; # showing that we can check here to see if all data was consumed&nbsp; &nbsp; &nbsp; &nbsp; # the except statement above effectively does the same thing&nbsp; &nbsp; &nbsp; &nbsp; try:&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; next(generator)&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; print("we didn't finish")&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; return&nbsp; &nbsp; &nbsp; &nbsp; except StopIteration:&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; print("we finished")@app.get("/infinite")async def infinite_stream():&nbsp; &nbsp; return StreamingResponse(streamer(infinite_generator()))@app.get("/finite")async def finite_stream():&nbsp; &nbsp; return StreamingResponse(streamer(finite_generator()))@app.get("/ainfinite")async def infinite_stream():&nbsp; &nbsp; return StreamingResponse(astreamer(infinite_generator()))@app.get("/afinite")async def finite_stream():&nbsp; &nbsp; return StreamingResponse(astreamer(finite_generator()))if __name__ == "__main__":&nbsp; &nbsp; uvicorn.run(app, host="0.0.0.0", port=8000)