Streaming a generated CSV with Flask

Question

I have this function for streaming text files:

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

I am working out how to stream a CSV in a similar manner. This page gives an example, but I wish to use the CSV module.

I can use StringIO and create a fresh "file" and CSV writer for each line, but it seems very inefficient. Is there a better way?

Answer 1

According to this answer how do I clear a stringio object? it is quicker to just create a new StringIO object for each line in the file than the method I use below. However if you still don't want to create new StringIO instances you can achieve what you want like this:

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

If you just want to stream back the results as they are created by csv.writer you can create a custom object implementing an interface the writer expects.

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