Flask异步端点中使用async generator流式输出报错求助

问题说明 #

Flask支持通过生成器实现内容流式输出，但在异步端点中直接传入async generator会触发错误：TypeError: 'async_generator' object is not iterable。

原代码 #

from flask import Flask, Response
import asyncio

app = Flask(__name__)

async def generate_data():
    for i in range(5):
        yield f"Data: {i}\n"
        await asyncio.sleep(1)

@app.route('/stream')
async def stream():
    return Response(generate_data(), mimetype='text/plain')

if __name__ == '__main__':
    app.run(debug=True, port=3030)

报错信息 #

╰─$ python test.py 
 * Serving Flask app 'test'
 * Debug mode: on
WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.
 * Running on http://127.0.0.1:3030
Press CTRL+C to quit
 * Restarting with stat
 * Debugger is active!
 * Debugger PIN: 899-865-629
Debugging middleware caught exception in streamed response at a point where response headers were already sent.
Traceback (most recent call last):
  File "/home/ashkan/Downloads/flask-async-test/venv/lib/python3.12/site-packages/werkzeug/wsgi.py", line 256, in __next__
    return self._next()
           ^^^^^^^^^^^^
  File "/home/ashkan/Downloads/flask-async-test/venv/lib/python3.12/site-packages/werkzeug/wrappers/response.py", line 32, in _iter_encoded
    for item in iterable:
                ^^^^^^^^^
TypeError: 'async_generator' object is not iterable
127.0.0.1 - - [29/Oct/2024 23:30:58] "GET /stream HTTP/1.1" 200 -

问题根源 #

Flask的Response依赖WSGI协议处理请求，而WSGI是同步模型，仅支持同步迭代器。async generator属于异步迭代器，无法被同步遍历，因此抛出类型错误。

解决方法 #

将异步生成器包装为同步迭代器，通过异步事件循环逐个获取输出内容，实现流式传输：

from flask import Flask, Response
import asyncio

app = Flask(__name__)

async def generate_data():
    for i in range(5):
        yield f"Data: {i}\n"
        await asyncio.sleep(1)

def sync_stream_wrapper():
    # 初始化异步事件循环
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    async_gen = generate_data()
    
    try:
        # 循环获取异步生成器的结果，转为同步输出
        while True:
            yield loop.run_until_complete(async_gen.__anext__())
    except StopAsyncIteration:
        # 捕获异步生成器结束信号，终止循环
        pass

@app.route('/stream')
async def stream():
    return Response(sync_stream_wrapper(), mimetype='text/plain')

if __name__ == '__main__':
    app.run(debug=True, port=3030)

修改后，访问/stream即可看到每隔1秒输出一行数据的流式效果。

内容的提问来源于stack exchange，提问作者Ashkan Arabi