support-formatter-api/support_formatter/routes/upload.py

import os
import pathlib
import uuid
from typing import List

import marshmallow as ma
from flask_smorest.fields import Upload
from werkzeug.datastructures import FileStorage

from ..app import Application
from ..config import APISettings
from ..logic.csv_processor import FileTypeInvalidError, process_csv
from ..utils import generate_error
from . import formatter_routes as blp

APP = Application.get_instance()

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in APISettings.ALLOWED_EXTENSIONS

class MultipartFormSchema(ma.Schema):
    username = ma.fields.String(required=True)
    friendcode = ma.fields.String(required=False)

class MultipartFileSchema(ma.Schema):
    servantdata = Upload()
    cedata = Upload()

@blp.route("/upload", methods=["POST"])
@blp.arguments(MultipartFormSchema, location="form")
@blp.arguments(MultipartFileSchema, location="files") 
@blp.response(200)
def upload_file(form: dict[str, str], files: dict[str, FileStorage]):
    filepaths: List[pathlib.Path] = []
    returndata = {
        "username": form["username"],
        "friendcode": form.get("friendcode", "")
    }

    for name, file in files.items():
        if name not in ("servantdata", "cedata"):
            return generate_error(406, message="Invalid form sent")

        filepath = APISettings.FILE_SAVE_DIRECTORY / f"{uuid.uuid4()}.csv"
        file.save(filepath)

        if os.stat(filepath).st_size < 1 or not allowed_file(file.filename):
            filepath.unlink(missing_ok=True)
            continue

        filepaths.append(filepath)

    if len(filepaths) == 0:
        return generate_error(406, message="No files provided")

    try:
        for f in filepaths:
            result = process_csv(f)
            returndata = returndata | result
            f.unlink(missing_ok=True)
    except FileTypeInvalidError:
        for f in filepaths:
            f.unlink(missing_ok=True)
        return generate_error(500, message="Error whilst parsing uploaded file - please contact Firq on Fate/Sacc Order")

    return returndata