Как провалидировать расширение и размер файла в FastAPI?

Рейтинг: 1Ответов: 0Опубликовано: 13.07.2023

Возникла необходимость провалидировать расширение и размер файла, однако, возникли сложности. Есть два варианта.

Вариант 1. Dependency injection

app = FastAPI()


class FileExtensionValidator:
    message = 'Extension “{extension}” not allowed. Allowed extensions are {allowed_extensions}'

    def __init__(self, allowed_extensions=None, message=None):
        if allowed_extensions is not None:
            allowed_extensions = [allowed_extension.lower() for allowed_extension in allowed_extensions]
        self.allowed_extensions = allowed_extensions
        if message is not None:
            self.message = message

    def __call__(self, file: UploadFile):
        extension = Path(file.filename).suffix[1:].lower()
        if self.allowed_extensions is not None and extension not in self.allowed_extensions:
            detail = self.message.format(extension=extension, allowed_extensions=', '.join(self.allowed_extensions))
            raise HTTPException(400, detail=detail)


class MaxFileSizeMBValidator:
    def __init__(self, max_mb: int):
        self.max_mb = max_mb

    def __call__(self, file: UploadFile):
        if file.size / 1024 / 1024 > self.max_mb:
            message = f'Maximum file size exceeded'
            raise HTTPException(400, detail=message)


@router.post(
    '/upload/',
    dependencies=[
        Depends(MaxFileSizeMBValidator(max_mb=1)),
        Depends(FileExtensionValidator(allowed_extensions=['.docx'])),
    ]
)
async def upload_file(file: UploadFile, user_id: int = Form()):
    return {'filename': file.filename}

Минусы данного варианта в том, что не возвращаются все ошибки валидации, тк как только будет выброшено исключение HTTPException, то FastAPI вернет ответ, а остальное не будет провалидировано. Так, например, если в запросе не будет передан user_id, то он не будет провалидирован и вернется ответ:

{
    "detail": "Extension “csv” not allowed. Allowed extensions are .docx"
}

Как видим, отработала самая первая объявленная зависимость. Я бы хотел получить ответ:

{
    "detail": [
        {
            "loc": [
                "body",
                "file"
            ],
            "msg": "Extension “csv” not allowed. Allowed extensions are .docx",
            "type": "value_error.missing"
        },
        {
            "loc": [
                "body",
                "user_id"
            ],
            "msg": "field required",
            "type": "value_error.missing"
        }
    ]
}

Можно это сделать?

Вариант 2. Переопределение UploadFile.

Посмотрел исходный код FastAPI, я узнал, что можно переопределить метод __get_validators__:

class CustomUploadFile(UploadFile):
    validators = [
        MaxFileSizeMBValidator(max_mb=10),
        FileExtensionValidator(allowed_extensions=['.docx']),
    ]

    @classmethod
    def __get_validators__(cls: Type["UploadFile"]) -> Iterable[Callable[..., Any]]:
        yield cls.validate

        for validator in cls.validators:
            yield validator

Но так придется каждый раз переопределять, а хочется как в Django как то так:

file = serializers.FileField(
    label='Файл для загрузки',
    validators=[
        FileExtensionValidator(allowed_extensions=['jpg', 'jpeg', 'png', 'pdf', 'bmp', 'tiff', 'tif', 'psd']),
        MaxFileSizeMBValidator(max_mb=50),
    ],
)

Можно как то так сделать?

Ответы

Ответов пока нет.