216.面试题-上传下载大文件django

import os

from pathlib import Path
from django.http import HttpResponse, StreamingHttpResponse

from tcp_client.settings import BASE_DIR

DEFAULT_DIR_PATH = Path(BASE_DIR).joinpath("libs")
MAX_SIZE = 1024 * 1024 * 2


# 自定义异常
class MyException(Exception):
    def __init__(self, code=1000, msg="ok"):
        self.code = code
        self.msg = msg


class UploadFileException(MyException):
    pass


class DownloadFileException(MyException):
    pass


class FileNotExistsException(MyException):
    pass


class FileOverMaxSizeException(MyException):
    pass


# 自定装饰器, 处理异常以及日志
def wrapper(func):
    def inner(*args):
        try:
            ret = func(*args)
        except MyException as e:
            ret = HttpResponse(e.msg)
        except Exception:
            import traceback
            ret = HttpResponse(traceback.format_exc())
        return ret
    return inner


@wrapper
def upload_file(request):
    """
    上传文件
    :param request.files: 文件参数
    :return: ret
    """
    files = request.FILES.getlist("files")
    if not files: raise UploadFileException(1001, "文件参数不能为空")
    for file in files:
        file_path = DEFAULT_DIR_PATH.joinpath(file.name)
        if not file_path.parent.exists():
            os.mkdir(file_path.parent)
        with open(file_path, "wb+") as f:
            for chunk in file.chunks():  # 每次写入一块, 节省内存
                f.write(chunk)
    ret = "ok"
    return HttpResponse(ret)


@wrapper
def download_file(request):
    """
    下载文件
    :param request.filename: 文件名称
    :return: ret
    """
    file_name = request.POST.get("filename")

    if not file_name: raise DownloadFileException(1001, "文件名称不能为空")
    file_path = DEFAULT_DIR_PATH.joinpath(file_name)
    if not file_path.exists(): raise FileNotExistsException(1001, "文件不存在")
    # if os.path.getsize(file_path) > MAX_SIZE: raise FileOverMaxSizeException(1001, "文件过大")

    # 使用生成器, 每次读取固定字节大小文件
    def file_iterator(file_path, chunk_size=1024):
        with open(file_path, "rb") as f:
            while True:
                c = f.read(chunk_size)
                if c:
                    yield c
                else:
                    break
    # django 自带的StreamingHttpResponse可以支持迭代器协议处理流文件
    ret = StreamingHttpResponse(file_iterator(file_path))
    # 将Content-Type设置为application/octet-stream流模式
    ret["Content-Type"] = "application/octet-stream"
    # Content-Disposition设置为attachment;filename={}, 浏览器会自动将保存到filename的文件中
    ret["Content-Disposition"] = "attachment;filename={}".format(file_name)
    return ret