mirror of https://github.com/tiangolo/fastapi.git
Enhance response model inference and improve error handling
- Updated the response model inference logic in "APIRoute" to check for "None" before validating against "BaseModel". - Refined the "infer_response_model_from_ast" function to handle nested function definitions and invalid dictionary keys, ensuring robust schema generation. - Added new test cases for edge scenarios to validate the inference behavior.
This commit is contained in:
parent
72b210209b
commit
40af1b90c3
|
|
@ -546,7 +546,7 @@ class APIRoute(routing.Route):
|
|||
else:
|
||||
response_model = return_annotation
|
||||
|
||||
if not lenient_issubclass(response_model, BaseModel):
|
||||
if response_model is not None and not lenient_issubclass(response_model, BaseModel):
|
||||
inferred = infer_response_model_from_ast(endpoint)
|
||||
if inferred:
|
||||
response_model = inferred
|
||||
|
|
|
|||
|
|
@ -366,11 +366,22 @@ def infer_response_model_from_ast(
|
|||
return None
|
||||
|
||||
return_stmt = None
|
||||
for node in ast.walk(func_def):
|
||||
|
||||
|
||||
nodes_to_visit = list(func_def.body)
|
||||
while nodes_to_visit:
|
||||
node = nodes_to_visit.pop(0)
|
||||
|
||||
if isinstance(node, ast.Return):
|
||||
return_stmt = node
|
||||
break
|
||||
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
||||
continue
|
||||
|
||||
for child in ast.iter_child_nodes(node):
|
||||
nodes_to_visit.append(child)
|
||||
|
||||
if not return_stmt:
|
||||
return None
|
||||
|
||||
|
|
@ -401,6 +412,10 @@ def infer_response_model_from_ast(
|
|||
for key, value in zip(dict_node.keys, dict_node.values):
|
||||
if not isinstance(key, ast.Constant):
|
||||
continue
|
||||
|
||||
if not isinstance(key.value, str):
|
||||
return None
|
||||
|
||||
field_name = key.value
|
||||
|
||||
field_type = _infer_type_from_ast(
|
||||
|
|
@ -418,4 +433,7 @@ def infer_response_model_from_ast(
|
|||
from pydantic import create_model
|
||||
|
||||
model_name = f"ResponseModel_{endpoint_function.__name__}"
|
||||
return create_model(model_name, **fields)
|
||||
try:
|
||||
return create_model(model_name, **fields)
|
||||
except Exception:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -1,5 +1,13 @@
|
|||
import sys
|
||||
import os
|
||||
import uvicorn
|
||||
|
||||
# Добавляем корень проекта в sys.path, чтобы Python видел пакет fastapi
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
from typing import Any, Dict, List
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
app = FastAPI()
|
||||
|
|
@ -71,6 +79,21 @@ async def get_local_variable() -> Dict[str, Any]:
|
|||
}
|
||||
return response_data
|
||||
|
||||
@app.get("/edge_cases/explicit_response")
|
||||
def get_explicit_response() -> JSONResponse:
|
||||
return JSONResponse({"should_not_be_inferred": True})
|
||||
|
||||
@app.get("/edge_cases/nested_function")
|
||||
async def get_nested_function() -> Dict[str, Any]:
|
||||
def inner_function():
|
||||
return {"inner": "value"}
|
||||
|
||||
return {"outer": "value"}
|
||||
|
||||
@app.get("/edge_cases/invalid_keys")
|
||||
def get_invalid_keys() -> Dict[Any, Any]:
|
||||
return {1: "value", "valid": "key"}
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
def test_openapi_schema_ast_inference():
|
||||
|
|
@ -111,4 +134,22 @@ def test_openapi_schema_ast_inference():
|
|||
|
||||
assert expr_props["calc_int"]["type"] == "integer"
|
||||
assert expr_props["calc_bool"]["type"] == "boolean"
|
||||
|
||||
|
||||
explicit_schema = paths["/edge_cases/explicit_response"]["get"]["responses"]["200"]["content"]["application/json"]["schema"]
|
||||
assert "$ref" not in explicit_schema
|
||||
|
||||
nested_schema = paths["/edge_cases/nested_function"]["get"]["responses"]["200"]["content"]["application/json"]["schema"]
|
||||
assert "$ref" in nested_schema
|
||||
nested_ref = nested_schema["$ref"].split("/")[-1]
|
||||
nested_props = schema["components"]["schemas"][nested_ref]["properties"]
|
||||
assert "outer" in nested_props
|
||||
assert "inner" not in nested_props
|
||||
|
||||
invalid_keys_schema = paths["/edge_cases/invalid_keys"]["get"]["responses"]["200"]["content"]["application/json"]["schema"]
|
||||
assert "$ref" not in invalid_keys_schema
|
||||
|
||||
if __name__ == "__main__":
|
||||
# test_openapi_schema_ast_inference()
|
||||
print("Запуск сервера для проверки Swagger UI...")
|
||||
print("Откройте в браузере: http://127.0.0.1:8000/docs")
|
||||
uvicorn.run(app, host="127.0.0.1", port=8000)
|
||||
|
|
|
|||
Loading…
Reference in New Issue