68 lines
2.8 KiB
Python
68 lines
2.8 KiB
Python
from typing import Any, Callable, Dict, Optional
|
||
from fastapi import FastAPI, Request, Response
|
||
from fastapi.responses import JSONResponse
|
||
import json
|
||
from starlette.middleware.base import BaseHTTPMiddleware
|
||
from starlette.types import ASGIApp
|
||
from app.schemas.response import StandardResponse, ErrorResponse
|
||
|
||
class ResponseMiddleware(BaseHTTPMiddleware):
|
||
"""
|
||
中间件：统一处理API响应格式
|
||
|
||
请求正确：{code:200, data:Any}
|
||
业务错误：{code:500, message:""}
|
||
"""
|
||
|
||
async def dispatch(self, request: Request, call_next: Callable) -> Response:
|
||
# 不需要处理的路径
|
||
exclude_paths = ["/docs", "/redoc", "/openapi.json"]
|
||
if any(request.url.path.startswith(path) for path in exclude_paths):
|
||
return await call_next(request)
|
||
|
||
response = await call_next(request)
|
||
|
||
# 如果是HTTPException，直接返回，不进行封装
|
||
if response.status_code >= 400:
|
||
return response
|
||
|
||
# 正常响应进行封装
|
||
if response.status_code < 300 and response.headers.get("content-type") == "application/json":
|
||
response_body = [chunk async for chunk in response.body_iterator]
|
||
response_body = b"".join(response_body)
|
||
|
||
if response_body:
|
||
try:
|
||
data = json.loads(response_body)
|
||
|
||
# 已经是统一格式的响应，不再封装
|
||
if isinstance(data, dict) and "code" in data and ("data" in data or "message" in data):
|
||
return Response(
|
||
content=response_body,
|
||
status_code=response.status_code,
|
||
headers=dict(response.headers),
|
||
media_type=response.media_type
|
||
)
|
||
|
||
# 封装为标准响应格式
|
||
result = StandardResponse(code=200, data=data).model_dump()
|
||
return JSONResponse(
|
||
content=result,
|
||
status_code=response.status_code,
|
||
headers=dict(response.headers)
|
||
)
|
||
|
||
except json.JSONDecodeError:
|
||
# 非JSON响应，直接返回
|
||
return Response(
|
||
content=response_body,
|
||
status_code=response.status_code,
|
||
headers=dict(response.headers),
|
||
media_type=response.media_type
|
||
)
|
||
|
||
return response
|
||
|
||
def add_response_middleware(app: FastAPI) -> None:
|
||
"""添加响应处理中间件到FastAPI应用"""
|
||
app.add_middleware(ResponseMiddleware) |