mirror
/
tiangolo.fastapi
mirror of https://github.com/tiangolo/fastapi
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

Merge branch 'feat/ast-response-inference' of https://github.com/g7AzaZLO/fastapi into feat/ast-response-inference

pull/14446/head
g7azazlo 8 months ago
parent
02493746e2 95d5b080a2
commit
6201559a1c
3 changed files with 98 additions and 78 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 19
      fastapi/routing.py
  2. 58
      fastapi/utils.py
  3. 99
      tests/test_ast_inference.py

19
fastapi/routing.py
View File

@ -545,16 +545,15 @@ class APIRoute(routing.Route):
response_model = None response_model = None
else: else:
response_model = return_annotation response_model = return_annotation
if (
response_model is None if (
or ( response_model is not None
not lenient_issubclass(response_model, BaseModel) and not lenient_issubclass(response_model, BaseModel)
and not dataclasses.is_dataclass(response_model) and not dataclasses.is_dataclass(response_model)
) ):
): inferred = infer_response_model_from_ast(endpoint)
inferred = infer_response_model_from_ast(endpoint) if inferred:
if inferred: response_model = inferred
response_model = inferred
self.response_model = response_model self.response_model = response_model
self.summary = summary self.summary = summary

58
fastapi/utils.py
View File

@ -1,7 +1,6 @@
import ast import ast
import inspect import inspect
import re import re
import warnings import warnings
from dataclasses import is_dataclass from dataclasses import is_dataclass
from typing import ( from typing import (
@ -277,12 +276,12 @@ def _infer_type_from_ast(
return List[Any] return List[Any]
first_type = _infer_type_from_ast(node.elts[0], func_def, context_name + "Item") first_type = _infer_type_from_ast(node.elts[0], func_def, context_name + "Item")
for elt in node.elts[1:]: for elt in node.elts[1:]:
current_type = _infer_type_from_ast(elt, func_def, context_name + "Item") current_type = _infer_type_from_ast(elt, func_def, context_name + "Item")
if current_type != first_type: if current_type != first_type:
return List[Any] return List[Any]
if first_type is not Any: if first_type is not Any:
return List[first_type] return List[first_type]
return List[Any] return List[Any]
@ -291,7 +290,7 @@ def _infer_type_from_ast(
left_type = _infer_type_from_ast(node.left, func_def, context_name) left_type = _infer_type_from_ast(node.left, func_def, context_name)
right_type = _infer_type_from_ast(node.right, func_def, context_name) right_type = _infer_type_from_ast(node.right, func_def, context_name)
if left_type == right_type and left_type in (int, float, str): if left_type == right_type and left_type in (int, float, str):
return left_type return left_type
if {left_type, right_type} == {int, float}: if {left_type, right_type} == {int, float}:
return float return float
@ -351,60 +350,63 @@ def infer_response_model_from_ast(
source = inspect.getsource(endpoint_function) source = inspect.getsource(endpoint_function)
except (OSError, TypeError): except (OSError, TypeError):
return None return None
source = inspect.cleandoc(source) source = inspect.cleandoc(source)
try: try:
tree = ast.parse(source) tree = ast.parse(source)
except SyntaxError: except SyntaxError:
return None return None
if not tree.body: if not tree.body:
return None return None
func_def = tree.body[0] func_def = tree.body[0]
if not isinstance(func_def, (ast.FunctionDef, ast.AsyncFunctionDef)): if not isinstance(func_def, (ast.FunctionDef, ast.AsyncFunctionDef)):
return None return None
return_stmt = None return_stmt = None
nodes_to_visit = list(func_def.body) nodes_to_visit = list(func_def.body)
while nodes_to_visit: while nodes_to_visit:
node = nodes_to_visit.pop(0) node = nodes_to_visit.pop(0)
if isinstance(node, ast.Return): if isinstance(node, ast.Return):
return_stmt = node return_stmt = node
break break
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
continue continue
for child in ast.iter_child_nodes(node): for child in ast.iter_child_nodes(node):
nodes_to_visit.append(child) nodes_to_visit.append(child)
if not return_stmt: if not return_stmt:
return None return None
returned_value = return_stmt.value returned_value = return_stmt.value
dict_node = None dict_node = None
if isinstance(returned_value, ast.Dict): if isinstance(returned_value, ast.Dict):
dict_node = returned_value dict_node = returned_value
elif isinstance(returned_value, ast.Name): elif isinstance(returned_value, ast.Name):
variable_name = returned_value.id variable_name = returned_value.id
# Find assignment # Find assignment
for node in func_def.body: for node in func_def.body:
if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == variable_name: if (
if isinstance(node.value, ast.Dict): isinstance(node, ast.AnnAssign)
dict_node = node.value and isinstance(node.target, ast.Name)
break and node.target.id == variable_name
):
if isinstance(node.value, ast.Dict):
dict_node = node.value
break
elif isinstance(node, ast.Assign): elif isinstance(node, ast.Assign):
for target in node.targets: for target in node.targets:
if isinstance(target, ast.Name) and target.id == variable_name: if isinstance(target, ast.Name) and target.id == variable_name:
if isinstance(node.value, ast.Dict): if isinstance(node.value, ast.Dict):
dict_node = node.value dict_node = node.value
break break
if not dict_node: if not dict_node:
return None return None
@ -412,16 +414,16 @@ def infer_response_model_from_ast(
for key, value in zip(dict_node.keys, dict_node.values): for key, value in zip(dict_node.keys, dict_node.values):
if not isinstance(key, ast.Constant): if not isinstance(key, ast.Constant):
continue continue
if not isinstance(key.value, str): if not isinstance(key.value, str):
return None return None
field_name = key.value field_name = key.value
field_type = _infer_type_from_ast( field_type = _infer_type_from_ast(
value, func_def, f"{endpoint_function.__name__}_{field_name}" value, func_def, f"{endpoint_function.__name__}_{field_name}"
) )
fields[field_name] = (field_type, ...) fields[field_name] = (field_type, ...)
if not fields: if not fields:

99
tests/test_ast_inference.py
View File

@ -1,10 +1,12 @@
from typing import Any, Dict, List from typing import Any, Dict
from fastapi import FastAPI from fastapi import FastAPI
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse
from fastapi.testclient import TestClient from fastapi.testclient import TestClient
app = FastAPI() app = FastAPI()
@app.get("/users/{user_id}") @app.get("/users/{user_id}")
async def get_user(user_id: int) -> Dict[str, Any]: async def get_user(user_id: int) -> Dict[str, Any]:
user: Dict[str, Any] = { user: Dict[str, Any] = {
@ -16,6 +18,7 @@ async def get_user(user_id: int) -> Dict[str, Any]:
} }
return user return user
@app.get("/orders/{order_id}") @app.get("/orders/{order_id}")
async def get_order_details(order_id: str) -> Dict[str, Any]: async def get_order_details(order_id: str) -> Dict[str, Any]:
order_data: Dict[str, Any] = { order_data: Dict[str, Any] = {
@ -40,49 +43,41 @@ async def get_order_details(order_id: str) -> Dict[str, Any]:
} }
return order_data return order_data
@app.get("/edge_cases/mixed_types") @app.get("/edge_cases/mixed_types")
async def get_mixed_types() -> Dict[str, Any]: async def get_mixed_types() -> Dict[str, Any]:
return { return {"mixed_list": [1, "two", 3.0], "description": "List starting with int"}
"mixed_list": [1, "two", 3.0],
"description": "List starting with int"
}
@app.get("/edge_cases/expressions") @app.get("/edge_cases/expressions")
async def get_expressions() -> Dict[str, Any]: async def get_expressions() -> Dict[str, Any]:
return { return {"calc_int": 10 + 5, "calc_str": "foo" + "bar", "calc_bool": 5 > 3}
"calc_int": 10 + 5,
"calc_str": "foo" + "bar",
"calc_bool": 5 > 3
}
@app.get("/edge_cases/empty_structures") @app.get("/edge_cases/empty_structures")
async def get_empty_structures() -> Dict[str, Any]: async def get_empty_structures() -> Dict[str, Any]:
return { return {"empty_list": [], "empty_dict": {}}
"empty_list": [],
"empty_dict": {}
}
@app.get("/edge_cases/local_variable") @app.get("/edge_cases/local_variable")
async def get_local_variable() -> Dict[str, Any]: async def get_local_variable() -> Dict[str, Any]:
response_data = { response_data = {"status": "ok", "nested": {"check": True}}
"status": "ok",
"nested": {
"check": True
}
}
return response_data return response_data
@app.get("/edge_cases/explicit_response") @app.get("/edge_cases/explicit_response")
def get_explicit_response() -> JSONResponse: def get_explicit_response() -> JSONResponse:
return JSONResponse({"should_not_be_inferred": True}) return JSONResponse({"should_not_be_inferred": True})
@app.get("/edge_cases/nested_function") @app.get("/edge_cases/nested_function")
async def get_nested_function() -> Dict[str, Any]: async def get_nested_function() -> Dict[str, Any]:
def inner_function(): def inner_function():
return {"inner": "value"} return {"inner": "value"}
return {"outer": "value"} return {"outer": "value"}
@app.get("/edge_cases/invalid_keys") @app.get("/edge_cases/invalid_keys")
def get_invalid_keys() -> Dict[Any, Any]: def get_invalid_keys() -> Dict[Any, Any]:
return {1: "value", "valid": "key"} return {1: "value", "valid": "key"}
@ -92,38 +87,44 @@ class FakeDB:
def get_user(self) -> Dict[str, Any]: def get_user(self) -> Dict[str, Any]:
return {"id": 1, "username": "db_user"} return {"id": 1, "username": "db_user"}
fake_db = FakeDB() fake_db = FakeDB()
@app.get("/db/direct_return") @app.get("/db/direct_return")
def get_db_direct() -> Dict[str, Any]: def get_db_direct() -> Dict[str, Any]:
return fake_db.get_user() return fake_db.get_user()
@app.get("/db/dict_construction") @app.get("/db/dict_construction")
def get_db_constructed() -> Dict[str, Any]: def get_db_constructed() -> Dict[str, Any]:
data = fake_db.get_user() data = fake_db.get_user()
return { return {"db_id": data["id"], "source": "database"}
"db_id": data["id"],
"source": "database"
}
client = TestClient(app) client = TestClient(app)
def test_openapi_schema_ast_inference(): def test_openapi_schema_ast_inference():
response = client.get("/openapi.json") response = client.get("/openapi.json")
assert response.status_code == 200 assert response.status_code == 200
schema = response.json() schema = response.json()
paths = schema["paths"] paths = schema["paths"]
user_schema = paths["/users/{user_id}"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] user_schema = paths["/users/{user_id}"]["get"]["responses"]["200"]["content"][
"application/json"
]["schema"]
assert "$ref" in user_schema assert "$ref" in user_schema
ref_name = user_schema["$ref"].split("/")[-1] ref_name = user_schema["$ref"].split("/")[-1]
user_props = schema["components"]["schemas"][ref_name]["properties"] user_props = schema["components"]["schemas"][ref_name]["properties"]
assert user_props["id"]["type"] == "integer" assert user_props["id"]["type"] == "integer"
assert user_props["username"]["type"] == "string" assert user_props["username"]["type"] == "string"
assert user_props["is_active"]["type"] == "boolean" assert user_props["is_active"]["type"] == "boolean"
order_schema = paths["/orders/{order_id}"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] order_schema = paths["/orders/{order_id}"]["get"]["responses"]["200"]["content"][
"application/json"
]["schema"]
assert "$ref" in order_schema assert "$ref" in order_schema
order_ref = order_schema["$ref"].split("/")[-1] order_ref = order_schema["$ref"].split("/")[-1]
order_props = schema["components"]["schemas"][order_ref]["properties"] order_props = schema["components"]["schemas"][order_ref]["properties"]
@ -135,39 +136,57 @@ def test_openapi_schema_ast_inference():
customer_prop = order_props["customer_info"] customer_prop = order_props["customer_info"]
assert "$ref" in customer_prop assert "$ref" in customer_prop
mixed_schema = paths["/edge_cases/mixed_types"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] mixed_schema = paths["/edge_cases/mixed_types"]["get"]["responses"]["200"][
"content"
]["application/json"]["schema"]
mixed_ref = mixed_schema["$ref"].split("/")[-1] mixed_ref = mixed_schema["$ref"].split("/")[-1]
mixed_props = schema["components"]["schemas"][mixed_ref]["properties"] mixed_props = schema["components"]["schemas"][mixed_ref]["properties"]
assert mixed_props["mixed_list"]["type"] == "array" assert mixed_props["mixed_list"]["type"] == "array"
expr_schema = paths["/edge_cases/expressions"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] expr_schema = paths["/edge_cases/expressions"]["get"]["responses"]["200"][
"content"
]["application/json"]["schema"]
expr_ref = expr_schema["$ref"].split("/")[-1] expr_ref = expr_schema["$ref"].split("/")[-1]
expr_props = schema["components"]["schemas"][expr_ref]["properties"] expr_props = schema["components"]["schemas"][expr_ref]["properties"]
assert expr_props["calc_int"]["type"] == "integer" assert expr_props["calc_int"]["type"] == "integer"
assert expr_props["calc_bool"]["type"] == "boolean" assert expr_props["calc_bool"]["type"] == "boolean"
explicit_schema = paths["/edge_cases/explicit_response"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] explicit_schema = paths["/edge_cases/explicit_response"]["get"]["responses"]["200"][
"content"
]["application/json"]["schema"]
assert "$ref" not in explicit_schema assert "$ref" not in explicit_schema
nested_schema = paths["/edge_cases/nested_function"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] nested_schema = paths["/edge_cases/nested_function"]["get"]["responses"]["200"][
"content"
]["application/json"]["schema"]
assert "$ref" in nested_schema assert "$ref" in nested_schema
nested_ref = nested_schema["$ref"].split("/")[-1] nested_ref = nested_schema["$ref"].split("/")[-1]
nested_props = schema["components"]["schemas"][nested_ref]["properties"] nested_props = schema["components"]["schemas"][nested_ref]["properties"]
assert "outer" in nested_props assert "outer" in nested_props
assert "inner" not in nested_props assert "inner" not in nested_props
invalid_keys_schema = paths["/edge_cases/invalid_keys"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] invalid_keys_schema = paths["/edge_cases/invalid_keys"]["get"]["responses"]["200"][
"content"
]["application/json"]["schema"]
assert "$ref" not in invalid_keys_schema assert "$ref" not in invalid_keys_schema
db_direct_schema = paths["/db/direct_return"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] db_direct_schema = paths["/db/direct_return"]["get"]["responses"]["200"]["content"][
"application/json"
]["schema"]
assert "$ref" not in db_direct_schema assert "$ref" not in db_direct_schema
db_constructed_schema = paths["/db/dict_construction"]["get"]["responses"]["200"]["content"]["application/json"]["schema"] db_constructed_schema = paths["/db/dict_construction"]["get"]["responses"]["200"][
"content"
]["application/json"]["schema"]
assert "$ref" in db_constructed_schema assert "$ref" in db_constructed_schema
db_constructed_ref = db_constructed_schema["$ref"].split("/")[-1] db_constructed_ref = db_constructed_schema["$ref"].split("/")[-1]
db_constructed_props = schema["components"]["schemas"][db_constructed_ref]["properties"] db_constructed_props = schema["components"]["schemas"][db_constructed_ref][
"properties"
assert db_constructed_props["source"]["type"] == "string" ]
assert "type" not in db_constructed_props["db_id"] or db_constructed_props["db_id"] == {}
assert db_constructed_props["source"]["type"] == "string"
assert (
"type" not in db_constructed_props["db_id"]
or db_constructed_props["db_id"] == {}
)

Write Preview
Loading…
Cancel
Save