diff --git a/sdk/python/src/datastar_py/quart.py b/sdk/python/src/datastar_py/quart.py index 78b8a0a92..f761f006b 100644 --- a/sdk/python/src/datastar_py/quart.py +++ b/sdk/python/src/datastar_py/quart.py @@ -1,6 +1,8 @@ from __future__ import annotations -from typing import Any +from collections.abc import AsyncIterable, Iterable +from inspect import isasyncgen, isgenerator +from typing import TYPE_CHECKING, Any from quart import make_response as _make_response from quart import request @@ -8,6 +10,9 @@ from . import _read_signals from .sse import SSE_HEADERS, ServerSentEventGenerator +if TYPE_CHECKING: + from quart.typing import HeadersValue, ResponseTypes, StatusCode + __all__ = [ "SSE_HEADERS", "ServerSentEventGenerator", @@ -16,9 +21,19 @@ ] -async def make_datastar_response(async_generator): - response = await _make_response(async_generator, SSE_HEADERS) - response.timeout = None +async def make_datastar_response( + response_content: str | Iterable | AsyncIterable | None, + status_or_headers: StatusCode | HeadersValue = None, + headers: HeadersValue = None, + /, +) -> ResponseTypes: + status = status_or_headers + if status_or_headers is not None and not isinstance(status_or_headers, int): + status, headers = None, status_or_headers + headers = {**SSE_HEADERS, **(headers or {})} + response = await _make_response(response_content, status, headers) + if isgenerator(response_content) or isasyncgen(response_content): + response.timeout = None return response