From 40408073a234dc2bdd917ff7f1fe4d2edf001032 Mon Sep 17 00:00:00 2001 From: xuyonghao <898441624@qq.com> Date: 星期三, 18 十二月 2024 14:10:29 +0800 Subject: [PATCH] Merge remote-tracking branch 'origin/master' --- app/api/excel.py | 53 +++++++++++++++++++++++++++++++++++------------------ 1 files changed, 35 insertions(+), 18 deletions(-) diff --git a/app/api/excel.py b/app/api/excel.py index 72488bd..a21ee3f 100644 --- a/app/api/excel.py +++ b/app/api/excel.py @@ -1,13 +1,7 @@ -from fastapi import APIRouter, File, UploadFile, Depends +from fastapi import APIRouter, File, UploadFile 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.utils.excelmerge.conformity import run_conformity import shutil import os @@ -17,6 +11,7 @@ ALLOWED_EXTENSIONS = {'xlsx'} EXCEL_FILES_PATH = 'data/output' SOURCE_FILES_PATH = 'data/source' + def allowed_file(filename): return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS @@ -56,8 +51,7 @@ 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(SOURCE_FILES_PATH, file.filename) with open(save_path, 'wb') as buffer: shutil.copyfileobj(file.file, buffer) save_path_list.append(save_path) @@ -70,20 +64,41 @@ @router.websocket("/ws/excel") async def ws_excel(websocket: WebSocket): await websocket.accept() + + create_dir_if_not_exists(SOURCE_FILES_PATH) + create_dir_if_not_exists(EXCEL_FILES_PATH) + 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 == "鏌ヨ鍚堝苟杩涘害": + if data == "\"鍚堝苟Excel\"": + run_excel = run_conformity() + files = os.listdir(EXCEL_FILES_PATH) + if run_excel: + first_file = files[0] + file_name = os.path.basename(first_file) + download_url = f"./api/document/download/{first_file}" + await websocket.send_json({ + "message": "鏂囨。鍚堝苟鎴愬姛锛�", + "type": "stream", + "files": [{ + "file_name": file_name, + "file_url": download_url + }] + }) + await websocket.send_json({ + "message": "鏂囨。鍚堝苟鎴愬姛锛�", + "type": "close", + }) + else: + await websocket.send_json({"error": "鍚堝苟澶辫触", "type": "stream", "files": []}) + elif data == "\"鏌ヨ鍚堝苟杩涘害\"": files = os.listdir(EXCEL_FILES_PATH) if not files: await websocket.send_json({"step_message": "姝e湪鍚堝苟涓�", "type": "stream", "files": []}) else: await websocket.send_json({"step_message": "鏂囨。鍚堝苟鎴愬姛锛�", "type": "stream", "files": []}) - elif data == "鑾峰彇鏂囦欢": + elif data == "\"鑾峰彇鏂囦欢\"": files = os.listdir(EXCEL_FILES_PATH) if not files: await websocket.send_json({"error": "鐩綍涓嬫病鏈夌敓鎴愮殑鏂囦欢", "type": "stream", "files": []}) @@ -100,9 +115,11 @@ }] }) 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}") @@ -113,4 +130,4 @@ except FileNotFoundError: raise HTTPException(status_code=404, detail="鏂囦欢涓嶅瓨鍦�") except Exception as e: - raise HTTPException(status_code=500, detail="鏈嶅姟鍣ㄩ敊璇�") \ No newline at end of file + raise HTTPException(status_code=500, detail="鏈嶅姟鍣ㄩ敊璇�") -- Gitblit v1.8.0