使用 Flask 流式传输生成的 CSV

Question

我有这个用于流式传输文本文件的功能：

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 并为每一行创建一个新的“文件”和 CSV 写入器，但这似乎效率很低。有没有更好的办法？

Answer 1

对贾斯汀现有的好答案略有改进。你可以利用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

Answer 2

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

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

Answer 3

根据这个答案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