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， 代碼也看起來神清氣爽，也不冗余了，頭也不疼不眩暈了，距離變禿又遠了一分。