背景

最近要用 Fastapi 框架简单实现一个文件下载接口，需要临时创建一个目录，将需要下载的文件打包为 zip 文件，返回给前端后，需要把这个 zip 文件删除。一开始我想到了两个思路：1. 维护一个定时器，定时去删除。2. 创建临时文件夹，接口返回后自动清理掉这个文件夹。

对比当然选择第二个方案，因为实现比较简单。

简单 demo 实现如下：

import os
import tempfile
import zipfile

import uvicorn
from fastapi import FastAPI
from fastapi.responses import FileResponse

app = FastAPI()


@app.get("/file/download")
async def file_download():
    base_dir = os.path.dirname(__file__)

    # Create temp dir
    tmp_dir_obj = tempfile.TemporaryDirectory(dir=base_dir)
    tmp_dirname = tmp_dir_obj.name

    # File to be zipped
    file_path = f'{base_dir}/data/test.txt'

    # Zip file
    zip_file_path = f'{tmp_dirname}/test.zip'
    with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as fp:
        fp.write(filename=file_path, arcname='test.txt')

    zip_filename = os.path.basename(zip_file_path)

    return FileResponse(path=zip_file_path, filename=zip_filename)


if __name__ == '__main__':
    uvicorn.run(app, port=8000)

可是当我使用 python 内置方法创建了临时文件夹，运行后却报以下错误：

image

找不到 zip 文件，因为 Fastapi 的请求返回是异步处理的，当前进程结束后，文件就被清理掉了，所以找不到这个文件对象。

于是乎我陷入了沉思。

因为实在不想用第一种实现（维护定时器还是觉得太麻烦了），我想了另一个思路：每次调这个下载接口，都会先去删除 zip 文件目录，这样可以保证服务器上最多只是保存一份 zip 文件。

但这种方式还是感觉很不优雅。因为上传文件时，要对保存上传文件的目录大小做限制，当超过设定值时要把最早的文件删除掉，以保证新的文件可以正常上传成功。上传文件目录和 zip 文件目录是在同一个目录下的，所以每次判断大小时，还要额外再判断下这个 zip 文件的大小。

解决

后面无意中发现了 Fastapi 是可以在某个请求结束后，后台执行任务的，主要借助了 starlette 这个包。具体实现如下：

import os
import zipfile

import uvicorn
from fastapi import FastAPI
from fastapi.responses import FileResponse
from starlette.background import BackgroundTask

app = FastAPI()


def rm_file(file_path):
    if os.path.exists(file_path):
        os.remove(file_path)


@app.get("/file/download")
async def file_download():
    base_dir = os.path.dirname(__file__)

    # File to be zipped
    file_path = f'{base_dir}/data/test.txt'

    # Zip file
    zip_file_path = f'{base_dir}/test.zip'
    with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as fp:
        fp.write(filename=file_path, arcname='test.txt')

    zip_filename = os.path.basename(zip_file_path)

    # Add background task
    task = BackgroundTask(rm_file, file_path=zip_file_path)

    return FileResponse(path=zip_file_path, filename=zip_filename, background=task)


if __name__ == '__main__':
    uvicorn.run(app, port=8000)

这样就可以完美清理掉临时生成的 zip 文件了。

参考

https://www.starlette.io/background/