使用 Flask 流式传输生成的 CSV

Streaming a generated CSV with Flask

我有流式传输文本文件的功能:

def txt_response(filename, iterator):
    if not filename.endswith('.txt'):
        filename += '.txt'
    filename = filename.format(date=str(datetime.date.today()).replace(' ', '_'))
    response = Response((_.encode('utf-8')+'\r\n' for _ in iterator), mimetype='text/txt')
    response.headers['Content-Disposition'] = 'attachment; filename={filename}'.format(filename=filename)
    return response

我正在研究如何以类似方式流式传输 CSV。 This page举个例子,我想用CSV模块

我可以使用 StringIO 并为每一行创建一个新的 "file" 和 CSV 编写器,但它似乎效率很低。有没有更好的方法?

根据这个答案 how do I clear a stringio object? 为文件中的每一行创建一个新的 StringIO 对象比我在下面使用的方法更快。但是,如果您仍然不想创建新的 StringIO 实例,您可以像这样实现您想要的:

import csv
import StringIO

from flask import Response


def iter_csv(data):
    line = StringIO.StringIO()
    writer = csv.writer(line)
    for csv_line in data:
        writer.writerow(csv_line)
        line.seek(0)
        yield line.read()
        line.truncate(0)
        line.seek(0)  # required for Python 3


def csv_response(data):
    response = Response(iter_csv(data), mimetype='text/csv')
    response.headers['Content-Disposition'] = 'attachment; filename=data.csv'
    return response

如果您只想流回 csv.writer 创建的结果,您可以创建一个实现作者期望的接口的自定义对象。

import csv

from flask import Response


class Line(object):
    def __init__(self):
        self._line = None
    def write(self, line):
        self._line = line
    def read(self):
        return self._line


def iter_csv(data):
    line = Line()
    writer = csv.writer(line)
    for csv_line in data:
        writer.writerow(csv_line)
        yield line.read()


def csv_response(data):
    response = Response(iter_csv(data), mimetype='text/csv')
    response.headers['Content-Disposition'] = 'attachment; filename=data.csv'
    return response

如果您要处理不想存储在内存中的大量数据,那么您可以使用 SpooledTemporaryFile。这将使用 StringIO 直到它达到 max_size 之后它将滚动到磁盘。

但是,如果您只想在创建结果时流式传输回结果,我会坚持推荐的答案。

对贾斯汀现有的出色答案略有改进。您可以利用 csv.writerow() returns the value returned by the underlying file's write call.

import csv
from flask import Response

class DummyWriter:
    def write(self, line):
        return line

def iter_csv(data):
    writer = csv.writer(DummyWriter())
    for row in data:
        yield writer.writerow(row)

def csv_response(data):
    response = Response(iter_csv(data), mimetype='text/csv')
    response.headers['Content-Disposition'] = 'attachment; filename=data.csv'
    return response