无知的 tonyseek

Yet Another Seeker

在 Flask 里产生流式响应

用过 Bottle [0] 的同学应该不会忘记它的流式响应 [1] ——在视图函数中使用 yield 关键字，让调用结果成为一个迭代器，那么 HTTP 客户端将会得到这个迭代器每次迭代的结果一部分，迭代器产生多少客户端收到多少，就像流一样。用这种方法在产生一些大的响应对象时（比如大文件下载），能有效地节约服务器内存。

运行以下代码并在浏览器访问 http://localhost:5000/stream

stream.py
 from time import sleep
 from bottle import route, run

 @route('/stream')
 def stream():
     yield 'START'
     sleep(3)
     yield 'MIDDLE'
     sleep(5)
     yield 'END'

 run()

就能看到 "START" 出现之后，过了三秒 "MIDDLE" 才出现，再过了五秒 "END" 出现，然后浏览器的小转盘停止旋转，响应结束。

这个特性其实不是 Bottle 本身的特性，而是 WSGI 中已经规定的支持 [2] 。WSGI 标准要求 WSGI Application 返回的就是迭代器，所以我们平时返回的“整块”响应其实都已被 Web 框架包装成只产生一个元素的迭代器。

比较纠结的是，在 Flask 却不能像 Bottle 一样这么便捷。如果在 Flask 的视图函数里面使用 yield ，那么将得到一个 TypeError 异常，异常提示信息说：

TypeError: 'generator' object is not callable

原来 Flask 要求视图函数返回的结果是可调用的。从 Flask 文档中得知，Flask 能够处理的响应是 werkzeug 中 Response 或其子类的实例。而我们平时使用视图函数，一般是返回一个字符串的，该字符串会被 Flask 包装成 Response 对象，也就是说这里给了一点点糖。具体做这件事的，是 flask.app.Flask 类中的 make_response 方法：

def make_response(self, rv):
    if rv is None:
        raise ValueError('View function did not return a response')
    if isinstance(rv, self.response_class):
        return rv
    if isinstance(rv, basestring):
        return self.response_class(rv)
    if isinstance(rv, tuple):
        return self.response_class(*rv)
    return self.response_class.force_type(rv, request.environ)

那么 Response 又是如何支持迭代器响应的呢？在 Flask 的邮件列表中有一个线索 [3] 提到：

yield is a keyword in Python that creates generators which are often the base of streaming functionality in frameworks. The underlaying Werkzeug system does support generators but Flask does not directly accept them. You can however use them with a trick:

该线索的作者给出了一个可行的方案，就是手动产生 Response 对象：

from flask import Response

@app.route('/foo')
def foo():
    def generate():
        yield 'first part'
        yield 'second part'
    return Response(generate(), direct_passthrough=True)

这么一来 Flask 就可以使用迭代器产生流式响应了。但是比起 Bottle 那么简洁的方式来说，这个 Response 的包装过程真是十分的罗嗦和丑陋啊。这个时候我们有两种选择，第一是编写一个装饰器：

from functools import update_wrapper
from flask import Response, current_app


class StreamView(object):
    """A decorator for flask view."""

    def __init__(self, view_function):
        self.view_function = view_function
        update_wrapper(self, self.view_function)

    def __call__(self, *args, **kwargs):
        return_value = self.view_function(*args, **kwargs)
        try:
            response = iter(return_value)
        except TypeError:
            # the return value is not iterable
            response = return_value
            current_app.logger.warning(
                "The stream view %r isn't iterable." % self)
        else:
            # the return value is iterable
            response = Response(return_value, direct_passthrough=True)
        return response


stream_view = StreamView

然后用它来装饰需要使用生成器（或返回迭代器）的视图函数：

@app.route('/stream')
@stream_view
def stream():
    yield 'START'
    sleep(3)
    yield 'MIDDLE'
    sleep(5)
    yield 'END'

非常简洁易用，个人很推荐。

第二种方法则适用于开发者觉得装饰器都太累赘的情况，这种情况我们可以给 Flask 的 make_response 打一个猴子补丁，让它像对字符串一样对迭代器进行包装。为了不要太丑陋，我用写 Flask Extension 的方法把这个猴子补丁写出来。

import types

class GeneratorView(object):
    """A Flask extension to support the generator view."""

    def __init__(self, app=None):
        self.app = None
        if app:
            self.init_app(app)

    def init_app(self, app):
        self.app = app
        # replace the ``make_response`` self.make_response = app.make_response
        app.make_response = self.make_stream_response

    def make_stream_response(self, rv):
        if isinstance(rv, types.GeneratorType):
            return self.app.response_class(
                rv, direct_passthrough=True)
        return self.make_response(rv)

然后像安装普通 Flask Extension 一样将这个扩展安装到 app 上，视图函数就彻底支持用生成器产生的流式响应的。值得注意的一点是在 make_stream_response 方法中我出于性能考虑使用 isinstance 而不是 try-except 来判别视图函数返回值类型，这样就无法对非生成器类型的迭代器生效了，这算是违背鸭子类型指导的后果吧（谁让你去检查它的 DNA 而不检查它会不会游泳）。如果需要返回其他类型迭代器，可能还是需要像前文的装饰器一样改成 try-except 模式。

最后附上一个挺好玩的例子，访问 http://localhost:5000/stream/5 就可以每隔一秒生成一个数字，直到第五秒结束。

from time import sleep
from flask import Flask

app = Flask(__name__)
generator_view = GeneratorView(app)

@app.route('/stream/<int:seconds>')
def stream(seconds):
    yield "<ul>"
    for i in range(int(seconds)):
        yield "<li>%d</li>" % (i + 1)
        sleep(1)
    yield "</ul>"

if __name__ == "__main__":
    app.run(debug=True)

这个例子用异步来实现更好，但我没有尝试过结合 gevent 和 Flask。Bottle 的文档倒是有例子。本文的缘起也正是那个例子 [4] 。

[0]Bottle
[1]Iterables and generators
[2]The application and framework side of WSGI
[3]Using yield
[4]Primer to Asynchronous Applications

Comments