flask动态生成csv

前言

最近在开发flask的一款小工具时，遇到了动态生成csv的问题。于是乎，翻阅了一些资料在这里做一下总结记录。

起因

最近想把某一组数据（比如下面这种）生成csv的形式，最初使用的方法时把数据先生成csv然后把文件下载下来，但是这种方法太过繁琐，于是决定试试动态生成csv。

user_data = [
    ['姓名', '性别', '年龄'],
    ['张三', '男', '18'],
    ['李四', '男', '19'],
    ['王五', '男', '20'],
    ['小花', '女', '18'],
    ['小丽', '女', '19'],
    ['小翠', '女', '20'],
]

经历了一番Google之后，找到了如下的方法。在这里记录分享一下。

生成文件

import os
import csv
from io import StringIO
from flask import Flask, Response, stream_with_context


app = Flask(__name__)

user_data = [
    ['姓名', '性别', '年龄'],
    ['张三', '男', '18'],
    ['李四', '男', '19'],
    ['王五', '男', '20'],
    ['小花', '女', '18'],
    ['小丽', '女', '19'],
    ['小翠', '女', '20'],
]


@app.route('/csv', methods=['GET'])
def download():
    def generate():
        data = StringIO()
        w = csv.writer(data)
        # writer
        for i in user_data:
            w.writerow(i)
            yield data.getvalue()
            data.seek(0)
            data.truncate(0)

    response = Response(stream_with_context(generate()), mimetype='text/csv')
    response.headers.set("Content-Disposition",
                         "attachment", filename="user.csv")
    return response


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

结果

可以看到成功生成了csv， 代码也看起来神清气爽，也不冗余了，头也不疼不眩晕了，距离变秃又远了一分。