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

✨ Implement support for JSON Lines and Streaming data with yield

pull/15022/head
Sebastián Ramírez 5 months ago
parent
e6ddf0c122
commit
07ab822c7c
3 changed files with 217 additions and 69 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. 32
      fastapi/dependencies/utils.py
  2. 51
      fastapi/openapi/utils.py
  3. 203
      fastapi/routing.py

32
fastapi/dependencies/utils.py
View File

@ -1,7 +1,17 @@
import dataclasses import dataclasses
import inspect import inspect
import sys import sys
from collections.abc import Callable, Mapping, Sequence from collections.abc import (
AsyncGenerator,
AsyncIterable,
AsyncIterator,
Callable,
Generator,
Iterable,
Iterator,
Mapping,
Sequence,
)
from contextlib import AsyncExitStack, contextmanager from contextlib import AsyncExitStack, contextmanager
from copy import copy, deepcopy from copy import copy, deepcopy
from dataclasses import dataclass from dataclasses import dataclass
@ -251,6 +261,26 @@ def get_typed_return_annotation(call: Callable[..., Any]) -> Any:
return get_typed_annotation(annotation, globalns) return get_typed_annotation(annotation, globalns)
_STREAM_ORIGINS = {
AsyncIterable,
AsyncIterator,
AsyncGenerator,
Iterable,
Iterator,
Generator,
}
def get_stream_item_type(annotation: Any) -> Any | None:
origin = get_origin(annotation)
if origin is not None and origin in _STREAM_ORIGINS:
type_args = get_args(annotation)
if type_args:
return type_args[0]
return Any
return None
def get_dependant( def get_dependant(
*, *,
path: str, path: str,

51
fastapi/openapi/utils.py
View File

@ -355,25 +355,40 @@ def get_openapi_path(
operation.setdefault("responses", {}).setdefault(status_code, {})[ operation.setdefault("responses", {}).setdefault(status_code, {})[
"description" "description"
] = route.response_description ] = route.response_description
if route_response_media_type and is_body_allowed_for_status_code( if is_body_allowed_for_status_code(route.status_code):
route.status_code # Check for JSONL streaming (generator endpoints)
): if route.is_json_stream:
response_schema = {"type": "string"} jsonl_content: dict[str, Any] = {}
if lenient_issubclass(current_response_class, JSONResponse): if route.stream_item_field:
if route.response_field: item_schema = get_schema_from_model_field(
response_schema = get_schema_from_model_field( field=route.stream_item_field,
field=route.response_field,
model_name_map=model_name_map, model_name_map=model_name_map,
field_mapping=field_mapping, field_mapping=field_mapping,
separate_input_output_schemas=separate_input_output_schemas, separate_input_output_schemas=separate_input_output_schemas,
) )
jsonl_content["itemSchema"] = item_schema
else: else:
response_schema = {} jsonl_content["itemSchema"] = {}
operation.setdefault("responses", {}).setdefault( operation.setdefault("responses", {}).setdefault(
status_code, {} status_code, {}
).setdefault("content", {}).setdefault(route_response_media_type, {})[ ).setdefault("content", {})["application/jsonl"] = jsonl_content
"schema" elif route_response_media_type:
] = response_schema response_schema = {"type": "string"}
if lenient_issubclass(current_response_class, JSONResponse):
if route.response_field:
response_schema = get_schema_from_model_field(
field=route.response_field,
model_name_map=model_name_map,
field_mapping=field_mapping,
separate_input_output_schemas=separate_input_output_schemas,
)
else:
response_schema = {}
operation.setdefault("responses", {}).setdefault(
status_code, {}
).setdefault("content", {}).setdefault(
route_response_media_type, {}
)["schema"] = response_schema
if route.responses: if route.responses:
operation_responses = operation.setdefault("responses", {}) operation_responses = operation.setdefault("responses", {})
for ( for (
@ -453,9 +468,9 @@ def get_fields_from_routes(
request_fields_from_routes: list[ModelField] = [] request_fields_from_routes: list[ModelField] = []
callback_flat_models: list[ModelField] = [] callback_flat_models: list[ModelField] = []
for route in routes: for route in routes:
if getattr(route, "include_in_schema", None) and isinstance( if not isinstance(route, routing.APIRoute):
route, routing.APIRoute continue
): if route.include_in_schema:
if route.body_field: if route.body_field:
assert isinstance(route.body_field, ModelField), ( assert isinstance(route.body_field, ModelField), (
"A request body must be a Pydantic Field" "A request body must be a Pydantic Field"
@ -465,6 +480,8 @@ def get_fields_from_routes(
responses_from_routes.append(route.response_field) responses_from_routes.append(route.response_field)
if route.response_fields: if route.response_fields:
responses_from_routes.extend(route.response_fields.values()) responses_from_routes.extend(route.response_fields.values())
if route.stream_item_field:
responses_from_routes.append(route.stream_item_field)
if route.callbacks: if route.callbacks:
callback_flat_models.extend(get_fields_from_routes(route.callbacks)) callback_flat_models.extend(get_fields_from_routes(route.callbacks))
params = get_flat_params(route.dependant) params = get_flat_params(route.dependant)

203
fastapi/routing.py
View File

@ -11,6 +11,7 @@ from collections.abc import (
Collection, Collection,
Coroutine, Coroutine,
Generator, Generator,
Iterator,
Mapping, Mapping,
Sequence, Sequence,
) )
@ -42,6 +43,7 @@ from fastapi.dependencies.utils import (
get_dependant, get_dependant,
get_flat_dependant, get_flat_dependant,
get_parameterless_sub_dependant, get_parameterless_sub_dependant,
get_stream_item_type,
get_typed_return_annotation, get_typed_return_annotation,
solve_dependencies, solve_dependencies,
) )
@ -66,7 +68,7 @@ from starlette._utils import is_async_callable
from starlette.concurrency import run_in_threadpool from starlette.concurrency import run_in_threadpool
from starlette.exceptions import HTTPException from starlette.exceptions import HTTPException
from starlette.requests import Request from starlette.requests import Request
from starlette.responses import JSONResponse, Response from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import ( from starlette.routing import (
BaseRoute, BaseRoute,
Match, Match,
@ -315,6 +317,24 @@ async def run_endpoint_function(
return await run_in_threadpool(dependant.call, **values) return await run_in_threadpool(dependant.call, **values)
def _build_response_args(
*, status_code: int | None, solved_result: Any
) -> dict[str, Any]:
response_args: dict[str, Any] = {
"background": solved_result.background_tasks,
}
# If status_code was set, use it, otherwise use the default from the
# response class, in the case of redirect it's 307
current_status_code = (
status_code if status_code else solved_result.response.status_code
)
if current_status_code is not None:
response_args["status_code"] = current_status_code
if solved_result.response.status_code:
response_args["status_code"] = solved_result.response.status_code
return response_args
def get_request_handler( def get_request_handler(
dependant: Dependant, dependant: Dependant,
body_field: ModelField | None = None, body_field: ModelField | None = None,
@ -330,6 +350,8 @@ def get_request_handler(
dependency_overrides_provider: Any | None = None, dependency_overrides_provider: Any | None = None,
embed_body_fields: bool = False, embed_body_fields: bool = False,
strict_content_type: bool | DefaultPlaceholder = Default(True), strict_content_type: bool | DefaultPlaceholder = Default(True),
stream_item_field: ModelField | None = None,
is_json_stream: bool = False,
) -> Callable[[Request], Coroutine[Any, Any, Response]]: ) -> Callable[[Request], Coroutine[Any, Any, Response]]:
assert dependant.call is not None, "dependant.call must be a function" assert dependant.call is not None, "dependant.call must be a function"
is_coroutine = dependant.is_coroutine_callable is_coroutine = dependant.is_coroutine_callable
@ -427,61 +449,115 @@ def get_request_handler(
embed_body_fields=embed_body_fields, embed_body_fields=embed_body_fields,
) )
errors = solved_result.errors errors = solved_result.errors
assert dependant.call # For types
if not errors: if not errors:
raw_response = await run_endpoint_function( if is_json_stream:
dependant=dependant, # Generator endpoint: stream as JSONL
values=solved_result.values, gen = dependant.call(**solved_result.values)
is_coroutine=is_coroutine,
) def _serialize_item(item: Any) -> bytes:
if isinstance(raw_response, Response): if stream_item_field:
if raw_response.background is None: value, errors = stream_item_field.validate(
raw_response.background = solved_result.background_tasks item, {}, loc=("response",)
response = raw_response )
else: if errors:
response_args: dict[str, Any] = { ctx = endpoint_ctx or EndpointContext()
"background": solved_result.background_tasks raise ResponseValidationError(
} errors=errors,
# If status_code was set, use it, otherwise use the default from the body=item,
# response class, in the case of redirect it's 307 endpoint_ctx=ctx,
current_status_code = ( )
status_code if status_code else solved_result.response.status_code line = stream_item_field.serialize_json(
value,
include=response_model_include,
exclude=response_model_exclude,
by_alias=response_model_by_alias,
exclude_unset=response_model_exclude_unset,
exclude_defaults=response_model_exclude_defaults,
exclude_none=response_model_exclude_none,
)
return line + b"\n"
else:
data = jsonable_encoder(item)
return json.dumps(data).encode("utf-8") + b"\n"
if dependant.is_async_gen_callable:
async def _async_stream_jsonl() -> AsyncIterator[bytes]:
async for item in gen:
yield _serialize_item(item)
stream_content: AsyncIterator[bytes] | Iterator[bytes] = (
_async_stream_jsonl()
)
else:
def _sync_stream_jsonl() -> Iterator[bytes]:
for item in gen:
yield _serialize_item(item)
stream_content = _sync_stream_jsonl()
response = StreamingResponse(
stream_content,
media_type="application/jsonl",
background=solved_result.background_tasks,
) )
if current_status_code is not None: response.headers.raw.extend(solved_result.response.headers.raw)
response_args["status_code"] = current_status_code elif dependant.is_async_gen_callable or dependant.is_gen_callable:
if solved_result.response.status_code: # Raw streaming with explicit response_class (e.g. StreamingResponse)
response_args["status_code"] = solved_result.response.status_code gen = dependant.call(**solved_result.values)
# Use the fast path (dump_json) when no custom response response_args = _build_response_args(
# class was set and a response field with a TypeAdapter status_code=status_code, solved_result=solved_result
# exists. Serializes directly to JSON bytes via Pydantic's
# Rust core, skipping the intermediate Python dict +
# json.dumps() step.
use_dump_json = response_field is not None and isinstance(
response_class, DefaultPlaceholder
) )
content = await serialize_response( response = actual_response_class(content=gen, **response_args)
field=response_field, response.headers.raw.extend(solved_result.response.headers.raw)
response_content=raw_response, else:
include=response_model_include, raw_response = await run_endpoint_function(
exclude=response_model_exclude, dependant=dependant,
by_alias=response_model_by_alias, values=solved_result.values,
exclude_unset=response_model_exclude_unset,
exclude_defaults=response_model_exclude_defaults,
exclude_none=response_model_exclude_none,
is_coroutine=is_coroutine, is_coroutine=is_coroutine,
endpoint_ctx=endpoint_ctx,
dump_json=use_dump_json,
) )
if use_dump_json: if isinstance(raw_response, Response):
response = Response( if raw_response.background is None:
content=content, raw_response.background = solved_result.background_tasks
media_type="application/json", response = raw_response
**response_args,
)
else: else:
response = actual_response_class(content, **response_args) response_args = _build_response_args(
if not is_body_allowed_for_status_code(response.status_code): status_code=status_code, solved_result=solved_result
response.body = b"" )
response.headers.raw.extend(solved_result.response.headers.raw) # Use the fast path (dump_json) when no custom response
# class was set and a response field with a TypeAdapter
# exists. Serializes directly to JSON bytes via Pydantic's
# Rust core, skipping the intermediate Python dict +
# json.dumps() step.
use_dump_json = response_field is not None and isinstance(
response_class, DefaultPlaceholder
)
content = await serialize_response(
field=response_field,
response_content=raw_response,
include=response_model_include,
exclude=response_model_exclude,
by_alias=response_model_by_alias,
exclude_unset=response_model_exclude_unset,
exclude_defaults=response_model_exclude_defaults,
exclude_none=response_model_exclude_none,
is_coroutine=is_coroutine,
endpoint_ctx=endpoint_ctx,
dump_json=use_dump_json,
)
if use_dump_json:
response = Response(
content=content,
media_type="application/json",
**response_args,
)
else:
response = actual_response_class(content, **response_args)
if not is_body_allowed_for_status_code(response.status_code):
response.body = b""
response.headers.raw.extend(solved_result.response.headers.raw)
if errors: if errors:
validation_error = RequestValidationError( validation_error = RequestValidationError(
errors, body=body, endpoint_ctx=endpoint_ctx errors, body=body, endpoint_ctx=endpoint_ctx
@ -609,12 +685,21 @@ class APIRoute(routing.Route):
) -> None: ) -> None:
self.path = path self.path = path
self.endpoint = endpoint self.endpoint = endpoint
self.stream_item_type: Any | None = None
if isinstance(response_model, DefaultPlaceholder): if isinstance(response_model, DefaultPlaceholder):
return_annotation = get_typed_return_annotation(endpoint) return_annotation = get_typed_return_annotation(endpoint)
if lenient_issubclass(return_annotation, Response): if lenient_issubclass(return_annotation, Response):
response_model = None response_model = None
else: else:
response_model = return_annotation stream_item = get_stream_item_type(return_annotation)
if stream_item is not None:
# Only extract item type for JSONL streaming when no
# explicit response_class (e.g. StreamingResponse) was set
if isinstance(response_class, DefaultPlaceholder):
self.stream_item_type = stream_item
response_model = None
else:
response_model = return_annotation
self.response_model = response_model self.response_model = response_model
self.summary = summary self.summary = summary
self.response_description = response_description self.response_description = response_description
@ -663,6 +748,15 @@ class APIRoute(routing.Route):
) )
else: else:
self.response_field = None # type: ignore self.response_field = None # type: ignore
if self.stream_item_type:
stream_item_name = "StreamItem_" + self.unique_id
self.stream_item_field: ModelField | None = create_model_field(
name=stream_item_name,
type_=self.stream_item_type,
mode="serialization",
)
else:
self.stream_item_field = None
self.dependencies = list(dependencies or []) self.dependencies = list(dependencies or [])
self.description = description or inspect.cleandoc(self.endpoint.__doc__ or "") self.description = description or inspect.cleandoc(self.endpoint.__doc__ or "")
# if a "form feed" character (page break) is found in the description text, # if a "form feed" character (page break) is found in the description text,
@ -704,6 +798,11 @@ class APIRoute(routing.Route):
name=self.unique_id, name=self.unique_id,
embed_body_fields=self._embed_body_fields, embed_body_fields=self._embed_body_fields,
) )
# Detect generator endpoints that should stream as JSONL
# (only when no explicit response_class like StreamingResponse is set)
self.is_json_stream = isinstance(response_class, DefaultPlaceholder) and (
self.dependant.is_async_gen_callable or self.dependant.is_gen_callable
)
self.app = request_response(self.get_route_handler()) self.app = request_response(self.get_route_handler())
def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any, Response]]: def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any, Response]]:
@ -722,6 +821,8 @@ class APIRoute(routing.Route):
dependency_overrides_provider=self.dependency_overrides_provider, dependency_overrides_provider=self.dependency_overrides_provider,
embed_body_fields=self._embed_body_fields, embed_body_fields=self._embed_body_fields,
strict_content_type=self.strict_content_type, strict_content_type=self.strict_content_type,
stream_item_field=self.stream_item_field,
is_json_stream=self.is_json_stream,
) )
def matches(self, scope: Scope) -> tuple[Match, Scope]: def matches(self, scope: Scope) -> tuple[Match, Scope]:

Write Preview
Loading…
Cancel
Save