Browse Source

Add Server-Sent Events (SSE) support

Add SSEResponse class and format_sse function for Server-Sent Events
support in FastAPI responses. Supports both async and sync generators,
event serialization, retry configuration, and automatic SSE formatting.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
pull/14977/head
Santosh Bhavani 5 months ago
parent
commit
064dece488
  1. 2
      fastapi/__init__.py
  2. 172
      fastapi/responses.py
  3. 201
      tests/test_sse_response.py

2
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 format_sse as format_sse
from .routing import APIRouter as APIRouter
from .websockets import WebSocket as WebSocket
from .websockets import WebSocketDisconnect as WebSocketDisconnect

172
fastapi/responses.py

@ -1,4 +1,6 @@
from typing import Any
from collections.abc import AsyncIterator
from collections.abc import Iterator
from fastapi.exceptions import FastAPIDeprecationWarning
from starlette.responses import FileResponse as FileResponse # noqa
@ -9,6 +11,176 @@ from starlette.responses import RedirectResponse as RedirectResponse # noqa
from starlette.responses import Response as Response # noqa
from starlette.responses import StreamingResponse as StreamingResponse # noqa
from typing_extensions import deprecated
import inspect
def _is_async(obj: Any) -> bool:
"""Check if an object is an async generator or async iterable."""
return inspect.isasyncgen(obj) or inspect.iscoroutine(obj)
class SSEResponse(StreamingResponse):
"""Server-Sent Events (SSE) response.
A response subclass that automatically formats events according to the SSE spec.
Supports automatic event formatting, retry configuration, and client disconnect handling.
Example:
```python
from fastapi import FastAPI
from fastapi.responses import SSEResponse
import asyncio
app = FastAPI()
@app.get("/events")
async def get_events():
async def event_generator():
for i in range(5):
yield {"event": "message", "data": f"Message {i}"}
await asyncio.sleep(1)
return SSEResponse(event_generator())
```
Or using the `event_serializer` parameter for simpler data:
```python
@app.get("/events")
async def get_events():
async def simple_generator():
for i in range(5):
yield f"Message {i}"
return SSEResponse(simple_generator(), event_serializer=lambda x: {"data": x})
```
"""
default_media_type = "text/event-stream"
def __init__(
self,
content: Any,
status_code: int = 200,
headers: dict[str, str] | None = None,
media_type: str | None = None,
background: Any = None,
encoding: str = "utf-8",
event_serializer: Any | None = None,
retry: int | None = None,
) -> None:
self.event_serializer = event_serializer
self.retry = retry
# Check if content is async
if _is_async(content):
# Create an SSE-formatted async iterator
async def sse_content() -> AsyncIterator[str]:
# First yield retry if specified
if retry is not None:
yield f"retry: {retry}\n\n"
# Then yield formatted events
async for item in content:
# Serialize the event data
if event_serializer is not None:
event_dict = event_serializer(item)
else:
event_dict = item
# Format as SSE
yield format_sse(event_dict)
super().__init__(
content=sse_content(),
status_code=status_code,
headers=headers,
media_type=media_type or self.default_media_type,
background=background,
encoding=encoding,
)
else:
# Create an SSE-formatted sync iterator
def sse_content() -> Iterator[str]:
# First yield retry if specified
if retry is not None:
yield f"retry: {retry}\n\n"
# Then yield formatted events
for item in content:
# Serialize the event data
if event_serializer is not None:
event_dict = event_serializer(item)
else:
event_dict = item
# Format as SSE
yield format_sse(event_dict)
super().__init__(
content=sse_content(),
status_code=status_code,
headers=headers,
media_type=media_type or self.default_media_type,
background=background,
encoding=encoding,
)
def format_sse(
data: str | dict[str, Any] = "",
event: str | None = None,
id: int | str | None = None,
retry: int | None = None,
) -> str:
"""Format data as a Server-Sent Event.
Args:
data: The event data. Can be a string or a dict with 'data', 'event', 'id', 'retry' keys.
event: The event type (optional).
id: The event ID (optional).
retry: The retry interval in milliseconds (optional).
Returns:
A properly formatted SSE string.
Example:
>>> format_sse(data="Hello")
'data: Hello\\n\\n'
>>> format_sse(data={"data": "Hello", "event": "message"})
'event: message\\ndata: Hello\\n\\n'
>>> format_sse("Hello", event="greeting", id="1", retry=5000)
'id: 1\\nevent: greeting\\nretry: 5000\\ndata: Hello\\n\\n'
"""
# Handle dict input
if isinstance(data, dict):
event = data.get("event", event)
id = data.get("id", id)
retry = data.get("retry", retry)
data = data.get("data", "")
lines = []
if id is not None:
lines.append(f"id: {id}")
if event is not None:
lines.append(f"event: {event}")
if retry is not None:
lines.append(f"retry: {retry}")
# Data can be multiline - each line should be prefixed with "data: "
if isinstance(data, str):
for line in data.split("\n"):
lines.append(f"data: {line}")
else:
lines.append(f"data: {data}")
# SSE events are separated by double newlines
return "\n".join(lines) + "\n\n"
try:
import ujson

201
tests/test_sse_response.py

@ -0,0 +1,201 @@
"""Tests for Server-Sent Events (SSE) response classes."""
import pytest
from fastapi import FastAPI
from fastapi.responses import SSEResponse, format_sse
from fastapi.testclient import TestClient
def test_format_sse_with_string_data():
"""Test format_sse with simple string data."""
result = format_sse("Hello")
assert result == "data: Hello\n\n"
def test_format_sse_with_dict_data():
"""Test format_sse with dict data."""
result = format_sse({"data": "Hello", "event": "message"})
assert result == "event: message\ndata: Hello\n\n"
def test_format_sse_with_all_params():
"""Test format_sse with all parameters."""
result = format_sse("Hello", event="greeting", id="1", retry=5000)
assert result == "id: 1\nevent: greeting\nretry: 5000\ndata: Hello\n\n"
def test_format_sse_multiline_data():
"""Test format_sse with multiline data."""
result = format_sse("Line1\nLine2\nLine3")
assert result == "data: Line1\ndata: Line2\ndata: Line3\n\n"
def test_format_sse_empty_data():
"""Test format_sse with empty data."""
result = format_sse("")
assert result == "data: \n\n"
def test_format_sse_with_event_only():
"""Test format_sse with only event type."""
result = format_sse(data="test", event="myevent")
assert result == "event: myevent\ndata: test\n\n"
def test_format_sse_with_retry():
"""Test format_sse with retry."""
result = format_sse(data="test", retry=3000)
assert result == "retry: 3000\ndata: test\n\n"
def test_format_sse_with_id():
"""Test format_sse with id."""
result = format_sse(data="test", id=42)
assert result == "id: 42\ndata: test\n\n"
# Integration tests with FastAPI
def test_sse_response_basic():
"""Test basic SSE response."""
app = FastAPI()
@app.get("/events")
async def get_events():
async def event_generator():
for i in range(3):
yield {"data": f"Message {i}"}
return SSEResponse(event_generator())
client = TestClient(app)
response = client.get("/events")
assert response.status_code == 200
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
assert "data: Message 0\n\n" in response.text
assert "data: Message 1\n\n" in response.text
assert "data: Message 2\n\n" in response.text
def test_sse_response_with_event_type():
"""Test SSE response with custom event type."""
app = FastAPI()
@app.get("/events")
async def get_events():
async def event_generator():
for i in range(2):
yield {"event": "notification", "data": f"Notice {i}"}
return SSEResponse(event_generator())
client = TestClient(app)
response = client.get("/events")
assert response.status_code == 200
assert "event: notification\ndata: Notice 0\n\n" in response.text
assert "event: notification\ndata: Notice 1\n\n" in response.text
def test_sse_response_with_retry():
"""Test SSE response with retry configuration."""
app = FastAPI()
@app.get("/events")
async def get_events():
async def event_generator():
for i in range(2):
yield {"data": f"Message {i}"}
return SSEResponse(event_generator(), retry=5000)
client = TestClient(app)
response = client.get("/events")
assert response.status_code == 200
# Retry is sent as first event
assert "retry: 5000\n\n" in response.text
def test_sse_response_with_event_serializer():
"""Test SSE response with custom event serializer."""
app = FastAPI()
@app.get("/events")
async def get_events():
async def simple_generator():
for i in range(3):
yield f"Message {i}"
def serializer(data: str) -> dict[str, str]:
return {"data": data}
return SSEResponse(simple_generator(), event_serializer=serializer)
client = TestClient(app)
response = client.get("/events")
assert response.status_code == 200
assert "data: Message 0\n\n" in response.text
assert "data: Message 1\n\n" in response.text
assert "data: Message 2\n\n" in response.text
def test_sse_response_with_sync_generator():
"""Test SSE response with synchronous generator."""
app = FastAPI()
@app.get("/events")
def get_events():
def sync_generator():
for i in range(3):
yield {"data": f"Sync {i}"}
return SSEResponse(sync_generator())
client = TestClient(app)
response = client.get("/events")
assert response.status_code == 200
assert "data: Sync 0\n\n" in response.text
assert "data: Sync 1\n\n" in response.text
assert "data: Sync 2\n\n" in response.text
def test_sse_response_custom_status_code():
"""Test SSE response with custom status code."""
app = FastAPI()
@app.get("/events")
async def get_events():
async def event_generator():
yield {"data": "test"}
return SSEResponse(event_generator(), status_code=201)
client = TestClient(app)
response = client.get("/events")
assert response.status_code == 201
def test_sse_response_custom_headers():
"""Test SSE response with custom headers."""
app = FastAPI()
@app.get("/events")
async def get_events():
async def event_generator():
yield {"data": "test"}
return SSEResponse(
event_generator(), headers={"X-Custom-Header": "custom-value"}
)
client = TestClient(app)
response = client.get("/events")
assert response.status_code == 200
assert response.headers["x-custom-header"] == "custom-value"
def test_sse_export():
"""Test that SSEResponse is exported from fastapi."""
from fastapi import SSEResponse
from fastapi import format_sse
assert SSEResponse is not None
assert format_sse is not None
Loading…
Cancel
Save