合并excel上传，合并，下载 修改

xuyonghao

2024-12-26 b40851d4f1d321782c17af6ed8f062113bf0ed31

 app/api/excel.py

@@ -1,13 +1,9 @@
from fastapi import APIRouter, File, UploadFile, Depends
from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, Depends
from fastapi.responses import JSONResponse, FileResponse
from fastapi.exceptions import HTTPException
from sqlalchemy.orm import Session
from starlette.websockets import WebSocket, WebSocketDisconnect
from werkzeug.utils import secure_filename
from app.api import get_current_user_websocket
from app.models.agent_model import AgentModel, AgentType
from app.models.base_model import get_db
from app.models.user_model import UserModel
from starlette.websockets import WebSocket

from app.api import get_current_user, get_current_user_websocket
from app.models import UserModel
from app.utils.excelmerge.conformity import run_conformity
import shutil
import os
@@ -18,17 +14,17 @@
EXCEL_FILES_PATH = 'data/output'
SOURCE_FILES_PATH = 'data/source'

def allowed_file(filename):

def allowed_file(filename: str) -> bool:
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def create_dir_if_not_exists(path):
def create_dir_if_not_exists(path: str):
    if not os.path.exists(path):
        os.makedirs(path)


# 清理函数
def clear_directory(path):
def clear_directory(path: str) -> dict:
    for filename in os.listdir(path):
        file_path = os.path.join(path, filename)
        try:
@@ -41,23 +37,30 @@
    return {"message": "目录已清空"}


def user_file_path(userid: str, path: str) -> str:
    return os.path.join(path, userid)


@router.post('/excel/upload')
async def upload_file(files: list[UploadFile] = File(...)):
async def upload_file(files: list[UploadFile] = File(...), current_user: UserModel = Depends(get_current_user)):
    user_id = str(current_user.id)
    if not any(file.filename for file in files):
        return JSONResponse(content={"error": "没有文件部分"}, status_code=400)
    if not user_id:
        return JSONResponse(content={"error": "缺少参数user_id"}, status_code=400)
    user_source = user_file_path(user_id, SOURCE_FILES_PATH)
    user_excel = EXCEL_FILES_PATH

    create_dir_if_not_exists(SOURCE_FILES_PATH)
    create_dir_if_not_exists(EXCEL_FILES_PATH)
    clear_directory(SOURCE_FILES_PATH)
    clear_directory(EXCEL_FILES_PATH)
    create_dir_if_not_exists(user_source)
    create_dir_if_not_exists(user_excel)
    clear_directory(user_source)

    save_path_list = []
    for file in files:
        if file.filename == '':
            return JSONResponse(content={"error": "没有选择文件"}, status_code=400)
        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            save_path = os.path.join(SOURCE_FILES_PATH, filename)
            save_path = os.path.join(user_source, file.filename)
            with open(save_path, 'wb') as buffer:
                shutil.copyfileobj(file.file, buffer)
            save_path_list.append(save_path)
@@ -68,55 +71,58 @@

# ws://localhost:9201/api/document/ws/excel
@router.websocket("/ws/excel")
async def ws_excel(websocket: WebSocket):
async def ws_excel(websocket: WebSocket, current_user: UserModel = Depends(get_current_user_websocket)):
    await websocket.accept()
    user_id = str(current_user.id)

    create_dir_if_not_exists(SOURCE_FILES_PATH)
    create_dir_if_not_exists(EXCEL_FILES_PATH)
    clear_directory(SOURCE_FILES_PATH)
    clear_directory(EXCEL_FILES_PATH)
    user_source = user_file_path(user_id, SOURCE_FILES_PATH)
    user_excel = EXCEL_FILES_PATH
    create_dir_if_not_exists(user_source)
    create_dir_if_not_exists(user_excel)

    while True:
        data = await websocket.receive_text()
        try:
            if data == "合并Excel文件":
                clear_directory(EXCEL_FILES_PATH)
                output_file_path = run_conformity()
                await websocket.send_json({"step_message": "开始合并", "type": "stream", "files": []})
            elif data == "查询合并进度":
                files = os.listdir(EXCEL_FILES_PATH)
                if not files:
                    await websocket.send_json({"step_message": "正在合并中", "type": "stream", "files": []})
                else:
                    await websocket.send_json({"step_message": "文档合并成功！", "type": "stream", "files": []})
            elif data == "获取文件":
                files = os.listdir(EXCEL_FILES_PATH)
                if not files:
                    await websocket.send_json({"error": "目录下没有生成的文件", "type": "stream", "files": []})
                else:
                    first_file = files[0]
                    file_name = os.path.basename(first_file)
                    file_url = f"./api/document/download/{first_file}"
            if data == "\"合并Excel\"":
                merge_file = run_conformity(user_source, user_excel)
                if merge_file is not None:

                    await websocket.send_json({
                        "step_message": "文档合并成功！",
                        "message": "文档合并成功！",
                        "type": "stream",
                        "files": [{
                            "file_name": file_name,
                            "file_url": file_url
                        }]
                        "file_name": f"{merge_file}.xlsx",
                        "download_url": f"./api/document/download/{merge_file}.xlsx"
                    })
                    await websocket.send_json({
                        "message": "文档合并成功！",
                        "type": "close",
                    })
                else:
                    await websocket.send_json({"error": "合并失败", "type": "stream", "files": []})
            else:
                await websocket.send_json({"error": "未知指令"})
                print(f"Received data: {data}")
                await websocket.send_json({"error": "未知指令", "data": str(data)})
        except Exception as e:
            await websocket.send_json({"error": str(e)})
            await websocket.close()


@router.get("/download/{filename}")
async def download_file(filename: str):
    try:
        return FileResponse(os.path.join(EXCEL_FILES_PATH, filename), filename=filename,
                            media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet')
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="文件不存在")
    except Exception as e:
        raise HTTPException(status_code=500, detail="服务器错误")
@router.get("/download/{file_full_name}")
async def download_file(background_tasks: BackgroundTasks, file_full_name: str):
    file_name = os.path.basename(file_full_name)
    user_excel = EXCEL_FILES_PATH
    file_path = os.path.join(user_excel, file_full_name)

    if not os.path.exists(file_path):
        return JSONResponse(content={"error": "文件不存在"}, status_code=404)

    def delete_file():
        try:
            os.unlink(file_path)
        except OSError as e:
            print(f"Deleting file error")

    # 待下载完成后删除生成的文件
    background_tasks.add_task(delete_file)
    return FileResponse(path=file_path, filename=file_name,
                        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")