使用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。 这个页面提供了一个示例，但我希望使用CSV模块。

我可以使用StringIO为每一行创建一个新的“文件”和CSV writer，但这似乎非常低效。有更好的方法吗？

- user964375

Answer 1

根据这个回答how do I clear a stringio object?，每一行的速度比我下面使用的方法快得多。但是，如果您仍然不想创建新的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

- Justin Fay

---

很好的答案。我还在for循环之前添加了writer.writeheader()和yield line.read()来添加标题。 - theannouncer

Answer 2

一份对Justin现有的出色答案的小改进。你可以利用这样一个事实，即csv.writerow()会返回基础文件的write调用所返回的值。 （链接）

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

- James

Answer 3

如果你正在处理大量数据，但不想将其存储在内存中，那么可以使用SpooledTemporaryFile。它会使用StringIO，直到达到max_size，然后会切换到磁盘上。

然而，如果你只是想流式返回已创建的结果，则建议使用推荐答案。

- Oswald