diff --git a/fastapi/__init__.py b/fastapi/__init__.py index 944fdd58f4..6ca6cd5cb7 100644 --- a/fastapi/__init__.py +++ b/fastapi/__init__.py @@ -20,6 +20,8 @@ from .param_functions import Query as Query from .param_functions import Security as Security from .requests import Request as Request from .responses import Response as Response +from .responses import SSEResponse as SSEResponse +from .responses import sse_content as sse_content from .routing import APIRouter as APIRouter from .websockets import WebSocket as WebSocket from .websockets import WebSocketDisconnect as WebSocketDisconnect diff --git a/fastapi/responses.py b/fastapi/responses.py index 5b1154c046..5a71d23490 100644 --- a/fastapi/responses.py +++ b/fastapi/responses.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, AsyncGenerator, Literal from fastapi.exceptions import FastAPIDeprecationWarning from starlette.responses import FileResponse as FileResponse # noqa @@ -10,6 +10,141 @@ from starlette.responses import Response as Response # noqa from starlette.responses import StreamingResponse as StreamingResponse # noqa from typing_extensions import deprecated + +class SSEResponse(StreamingResponse): + """Streaming response for Server-Sent Events (SSE). + + This response class automatically formats events according to the SSE + specification and provides convenient methods for sending events. + + Read more about Server-Sent Events in the + [FastAPI docs](https://fastapi.tiangolo.com/advanced/streaming-responses/#server-sent-events). + + Example: + ```python + from fastapi import FastAPI + from fastapi.responses import SSEResponse + + app = FastAPI() + + @app.get("/events") + def events(): + def event_generator(): + yield {"event": "message", "data": "Hello"} + yield {"event": "message", "data": "World"} + + return SSEResponse(event_generator()) + ``` + """ + + default_media_type: str = "text/event-stream" + + def __init__( + self, + content: AsyncGenerator[str, None] | AsyncGenerator[dict, None], + status_code: int = 200, + headers: dict[str, str] | None = None, + media_type: str | None = None, + background: Any = None, + retry: int | None = None, + ) -> None: + self._retry = retry + super().__init__( + content=self._wrap_content(content), + status_code=status_code, + headers=headers, + media_type=media_type or self.default_media_type, + background=background, + ) + + async def _wrap_content( + self, + content: AsyncGenerator[str, None] | AsyncGenerator[dict, None], + ) -> AsyncGenerator[bytes, None]: + """Wrap the content to format SSE events.""" + async for item in content: + if isinstance(item, dict): + yield self._format_event(item).encode("utf-8") + else: + yield item.encode("utf-8") + + def _format_event(self, event_data: dict[str, Any]) -> str: + """Format a dictionary as an SSE event.""" + lines = [] + + if "event" in event_data: + lines.append(f"event: {event_data['event']}") + + if "data" in event_data: + data = event_data["data"] + if isinstance(data, (list, dict)): + import json + + data = json.dumps(data) + for line in str(data).split("\n"): + lines.append(f"data: {line}") + + if "id" in event_data: + lines.append(f"id: {event_data['id']}") + + if "retry" in event_data: + lines.append(f"retry: {event_data['retry']}") + elif self._retry is not None: + lines.append(f"retry: {self._retry}") + + lines.append("") + return "\n".join(lines) + + +async def sse_content( + iterator: AsyncGenerator[dict[str, Any], None], +) -> AsyncGenerator[str, None]: + """Helper to create SSE-formatted content from a dictionary iterator. + + This is a convenience function that yields pre-formatted SSE strings, + useful when you want full control over the event formatting. + + Example: + ```python + from fastapi import FastAPI + from fastapi.responses import StreamingResponse, sse_content + + app = FastAPI() + + @app.get("/events") + async def events(): + async def generator(): + yield {"event": "message", "data": "Hello"} + yield {"data": "World"} + + return StreamingResponse(sse_content(generator()), media_type="text/event-stream") + ``` + """ + async for event in iterator: + lines = [] + + if "event" in event: + lines.append(f"event: {event['event']}") + + if "data" in event: + data = event["data"] + if isinstance(data, (list, dict)): + import json + + data = json.dumps(data) + for line in str(data).split("\n"): + lines.append(f"data: {line}") + + if "id" in event: + lines.append(f"id: {event['id']}") + + if "retry" in event: + lines.append(f"retry: {event['retry']}") + + lines.append("") + yield "\n".join(lines) + + try: import ujson except ImportError: # pragma: nocover diff --git a/tests/test_sse_response.py b/tests/test_sse_response.py new file mode 100644 index 0000000000..08724e9fdc --- /dev/null +++ b/tests/test_sse_response.py @@ -0,0 +1,224 @@ +"""Tests for Server-Sent Events (SSE) response.""" +from collections.abc import AsyncGenerator +from typing import Any + +import pytest +from fastapi import FastAPI +from fastapi.responses import SSEResponse, sse_content +from fastapi.testclient import TestClient + + +def test_sse_response_basic(): + """Test basic SSE response with event data.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + yield {"event": "message", "data": "Hello"} + yield {"data": "World"} + + return SSEResponse(generator()) + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + assert "text/event-stream" in response.headers["content-type"] + # Check that the SSE format is correct + content = response.text + assert "event: message" in content + assert "data: Hello" in content + assert "data: World" in content + + +def test_sse_response_with_retry(): + """Test SSE response with retry configuration.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + yield {"data": "Hello"} + + return SSEResponse(generator(), retry=5000) + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + assert "retry: 5000" in response.text + + +def test_sse_response_with_event_id(): + """Test SSE response with event ID.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + yield {"event": "update", "data": "test", "id": "1"} + + return SSEResponse(generator()) + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + assert "event: update" in response.text + assert "data: test" in response.text + assert "id: 1" in response.text + + +def test_sse_response_string_content(): + """Test SSE response with pre-formatted string content.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[str, None]: + yield "data: Hello\n\n" + + return SSEResponse(generator()) + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + assert "data: Hello" in response.text + + +def test_sse_response_json_data(): + """Test SSE response with JSON data.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + yield {"data": {"message": "hello", "count": 42}} + + return SSEResponse(generator()) + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + # JSON data should be serialized + assert '"message": "hello"' in response.text + assert '"count": 42' in response.text + + +def test_sse_response_with_retry_in_event(): + """Test SSE response retry in event data overrides default.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + # Event-level retry should override constructor retry + yield {"data": "Hello", "retry": 3000} + + return SSEResponse(generator(), retry=5000) + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + assert "retry: 3000" in response.text + # The default retry should not appear + lines = response.text.split("\n") + retry_count = sum(1 for line in lines if line.startswith("retry:")) + assert retry_count == 1 + + +def test_sse_content_helper(): + """Test the sse_content helper function.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + yield {"event": "message", "data": "Hello"} + yield {"data": "World"} + + from starlette.responses import StreamingResponse + + return StreamingResponse(sse_content(generator()), media_type="text/event-stream") + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + assert "text/event-stream" in response.headers["content-type"] + content = response.text + assert "event: message" in content + assert "data: Hello" in content + assert "data: World" in content + + +def test_sse_response_custom_media_type(): + """Test SSE response with custom media type.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + yield {"data": "Hello"} + + return SSEResponse(generator(), media_type="text/event-stream; charset=utf-8") + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + assert "text/event-stream" in response.headers["content-type"] + + +def test_sse_response_data_newlines(): + """Test SSE response with multiline data.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + yield {"data": "line1\nline2\nline3"} + + return SSEResponse(generator()) + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + # Each line should be prefixed with "data: " + assert "data: line1" in response.text + assert "data: line2" in response.text + assert "data: line3" in response.text + + +def test_sse_response_multiple_events(): + """Test SSE response with multiple events.""" + app = FastAPI() + + @app.get("/events") + def events(): + async def generator() -> AsyncGenerator[dict[str, Any], None]: + yield {"event": "first", "data": "data1"} + yield {"event": "second", "data": "data2", "id": "2"} + yield {"data": "data3"} + + return SSEResponse(generator()) + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/events") + + assert response.status_code == 200 + content = response.text + # First event + assert "event: first" in content + assert "data: data1" in content + # Second event + assert "event: second" in content + assert "data: data2" in content + assert "id: 2" in content + # Third event (no event type) + assert content.count("data: data3") == 1