StreamingResponse and pydantic.BaseModel?

[jamesbraza]

Previously, my endpoint directly returned a list[pydantic.BaseModel]. Now, I am trying to move to StreamingResponse with an AsyncIterable yielding pydantic.BaseModels.

How can I get StreamingResponse to work with my iterable of pydantic.BaseModel? Do I need to stringify the pydantic.BaseModel to get it to work with StreamingResponse?

[jamesbraza]

Note that #419 is related.

You can do something like this (pydantic==2.5.2, starlette==0.27.0):

from collections.abc import AsyncIterable, Iterable

from pydantic import BaseModel
from starlette.background import BackgroundTask
from starlette.concurrency import iterate_in_threadpool
from starlette.responses import JSONResponse, StreamingResponse


class JSONStreamingResponse(StreamingResponse, JSONResponse):
    """StreamingResponse that also render with JSON."""

    def __init__(
        self,
        content: Iterable | AsyncIterable,
        status_code: int = 200,
        headers: dict[str, str] | None = None,
        media_type: str | None = None,
        background: BackgroundTask | None = None,
    ) -> None:
        if isinstance(content, AsyncIterable):
            self._content_iterable: AsyncIterable = content
        else:
            self._content_iterable = iterate_in_threadpool(content)

        async def body_iterator() -> AsyncIterable[bytes]:
            async for content_ in self._content_iterable:
                if isinstance(content_, BaseModel):
                    content_ = content_.model_dump()
                yield self.render(content_)

        self.body_iterator = body_iterator()
        self.status_code = status_code
        if media_type is not None:
            self.media_type = media_type
        self.background = background
        self.init_headers(headers)