Skip to content

Commit b3a68f0

Browse files
gazpachokinggazpachoking
committed
Add sse_generator decorator to reduce boilerplate for SSE responses
1 parent 3379e82 commit b3a68f0

File tree

7 files changed

+102
-23
lines changed

7 files changed

+102
-23
lines changed

‎examples/python/fastapi/app.py‎

Lines changed: 36 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
import asyncio
22
from datetime import datetime
33

4+
import uvicorn
45
from fastapi import FastAPI
56
from fastapi.responses import HTMLResponse
67

7-
from datastar_py.fastapi import DatastarStreamingResponse
8+
from datastar_py.fastapi import DatastarStreamingResponse, ServerSentEventGenerator
9+
from datastar_py.starlette import sse_generator
810

911
app = FastAPI()
1012

@@ -15,7 +17,7 @@
1517
<head>
1618
<title>DATASTAR on FastAPI</title>
1719
<meta name="viewport" content="width=device-width, initial-scale=1.0"/>
18-
<script type="module" src="https://cdn.jsdelivr.net/gh/starfederation/datastar/bundles/datastar.js"></script>
20+
<script type="module" src="https://cdn.jsdelivr.net/gh/starfederation/datastar@v1.0.0-beta.11/bundles/datastar.js"></script>
1921
<style>
2022
html, body { height: 100%; width: 100%; }
2123
body { background-image: linear-gradient(to right bottom, oklch(0.424958 0.052808 253.972015), oklch(0.189627 0.038744 264.832977)); }
@@ -29,7 +31,7 @@
2931
<div class="container">
3032
<div
3133
class="time"
32-
data-on-load="@get('/updates')"
34+
data-on-load="@get('/updates3')"
3335
>
3436
Current time from fragment: <span id="currentTime">CURRENT_TIME</span>
3537
</div>
@@ -53,14 +55,43 @@ async def read_root():
5355

5456
async def time_updates():
5557
while True:
56-
yield DatastarStreamingResponse.merge_fragments(
58+
yield ServerSentEventGenerator.merge_fragments(
5759
[f"""<span id="currentTime">{datetime.now().isoformat()}"""]
5860
)
5961
await asyncio.sleep(1)
60-
yield DatastarStreamingResponse.merge_signals({"currentTime": f"{datetime.now().isoformat()}"})
62+
yield ServerSentEventGenerator.merge_signals({"currentTime": f"{datetime.now().isoformat()}"})
6163
await asyncio.sleep(1)
6264

6365

6466
@app.get("/updates")
6567
async def updates():
6668
return DatastarStreamingResponse(time_updates())
69+
70+
71+
# This is identical to the /updates endpoint, but uses a convenience decorator
72+
@app.get("/updates2")
73+
@sse_generator
74+
async def updates2():
75+
while True:
76+
yield ServerSentEventGenerator.merge_fragments(
77+
[f"""<span id="currentTime">{datetime.now().isoformat()}"""]
78+
)
79+
await asyncio.sleep(1)
80+
yield ServerSentEventGenerator.merge_signals({"currentTime": f"{datetime.now().isoformat()}"})
81+
await asyncio.sleep(1)
82+
83+
84+
# This is also identical, but yielding lists automaticall calls merge_fragments
85+
# and dicts automaticall calls merge_signals
86+
@app.get("/updates3")
87+
@sse_generator
88+
async def updates3():
89+
while True:
90+
yield [f"""<span id="currentTime">{datetime.now().isoformat()}"""]
91+
await asyncio.sleep(1)
92+
yield {"currentTime": f"{datetime.now().isoformat()}"}
93+
await asyncio.sleep(1)
94+
95+
96+
if __name__ == '__main__':
97+
uvicorn.run(app)

‎examples/python/quart/app.py‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
from quart import Quart
55

6-
from datastar_py.quart import make_datastar_response, ServerSentEventGenerator
6+
from datastar_py.quart import make_datastar_response, ServerSentEventGenerator, sse_generator
77

88
app = Quart(__name__)
99

@@ -13,7 +13,7 @@
1313
<head>
1414
<title>DATASTAR on Quart</title>
1515
<meta name="viewport" content="width=device-width, initial-scale=1.0"/>
16-
<script type="module" src="https://cdn.jsdelivr.net/gh/starfederation/datastar/bundles/datastar.js"></script>
16+
<script type="module" src="https://cdn.jsdelivr.net/gh/starfederation/datastar@v1.0.0-beta.11/bundles/datastar.js"></script>
1717
<style>
1818
html, body { height: 100%; width: 100%; }
1919
body { background-image: linear-gradient(to right bottom, oklch(0.424958 0.052808 253.972015), oklch(0.189627 0.038744 264.832977)); }
@@ -43,23 +43,18 @@
4343

4444

4545
@app.route("/updates")
46+
@sse_generator
4647
async def updates():
47-
async def time_updates():
48-
while True:
49-
yield ServerSentEventGenerator.merge_fragments(
50-
[f"""<span id="currentTime">{datetime.now().isoformat()}"""]
51-
)
52-
await asyncio.sleep(1)
53-
yield ServerSentEventGenerator.merge_signals({"currentTime": f"{datetime.now().isoformat()}"})
54-
await asyncio.sleep(1)
55-
56-
response = await make_datastar_response(time_updates())
57-
return response
48+
while True:
49+
yield [f"""<span id="currentTime">{datetime.now().isoformat()}"""]
50+
await asyncio.sleep(1)
51+
yield ServerSentEventGenerator.merge_signals({"currentTime": f"{datetime.now().isoformat()}"})
52+
await asyncio.sleep(1)
5853

5954

6055
@app.route("/")
6156
async def hello():
6257
return HTML.replace("CURRENT_TIME", f"{datetime.isoformat(datetime.now())}")
6358

6459

65-
# app.run()
60+
app.run()

‎sdk/python/src/datastar_py/django.py‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,26 @@
1+
import typing
12
from functools import wraps
23

34
from django.http import StreamingHttpResponse as _StreamingHttpResponse
45

5-
from .sse import SSE_HEADERS, ServerSentEventGenerator
6+
from .sse import SSE_HEADERS, ServerSentEventGenerator, _async_map, _wrap_event
67

78

89
class DatastarStreamingHttpResponse(_StreamingHttpResponse, ServerSentEventGenerator):
910
@wraps(_StreamingHttpResponse.__init__)
1011
def __init__(self, *args, **kwargs):
1112
kwargs["headers"] = {**SSE_HEADERS, **kwargs.get("headers", {})}
1213
super().__init__(*args, **kwargs)
14+
15+
16+
def sse_generator(generator_func):
17+
@wraps(generator_func)
18+
def _wrapper(*args, **kwargs):
19+
content = generator_func(*args, **kwargs)
20+
if isinstance(content, typing.AsyncIterable):
21+
content = _async_map(_wrap_event, content)
22+
else:
23+
content = map(_wrap_event, content)
24+
25+
return DatastarStreamingHttpResponse(content)
26+
return _wrapper

‎sdk/python/src/datastar_py/fastapi.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
from .starlette import DatastarStreamingResponse, ServerSentEventGenerator
1+
from .starlette import DatastarStreamingResponse, ServerSentEventGenerator, sse_generator

‎sdk/python/src/datastar_py/quart.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,20 @@
1+
from functools import wraps
2+
13
from quart import make_response as _make_response
24

3-
from .sse import ServerSentEventGenerator, SSE_HEADERS
5+
from .sse import ServerSentEventGenerator, SSE_HEADERS, _async_map, _wrap_event
46

57

68
async def make_datastar_response(async_generator):
79
response = await _make_response(async_generator, SSE_HEADERS)
810
response.timeout = None
911
return response
12+
13+
14+
def sse_generator(generator_func):
15+
@wraps(generator_func)
16+
async def _wrapper(*args, **kwargs):
17+
content = _async_map(_wrap_event, generator_func(*args, **kwargs))
18+
19+
return await make_datastar_response(content)
20+
return _wrapper

‎sdk/python/src/datastar_py/sse.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,3 +162,17 @@ def execute_script(
162162
event_id,
163163
retry_duration,
164164
)
165+
166+
167+
def _wrap_event(event):
168+
if isinstance(event, list):
169+
return ServerSentEventGenerator.merge_fragments(event)
170+
elif isinstance(event, dict):
171+
return ServerSentEventGenerator.merge_signals(event)
172+
else:
173+
return event
174+
175+
176+
async def _async_map(func, async_iter):
177+
async for item in async_iter:
178+
yield func(item)

‎sdk/python/src/datastar_py/starlette.py‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,26 @@
1+
import typing
12
from functools import wraps
23

34
from starlette.responses import StreamingResponse as _StreamingResponse
45

5-
from .sse import SSE_HEADERS, ServerSentEventGenerator
6+
from .sse import SSE_HEADERS, ServerSentEventGenerator, _wrap_event, _async_map
67

78

89
class DatastarStreamingResponse(_StreamingResponse, ServerSentEventGenerator):
910
@wraps(_StreamingResponse.__init__)
1011
def __init__(self, *args, **kwargs):
1112
kwargs["headers"] = {**SSE_HEADERS, **kwargs.get("headers", {})}
1213
super().__init__(*args, **kwargs)
14+
15+
16+
def sse_generator(generator_func):
17+
@wraps(generator_func)
18+
def _wrapper(*args, **kwargs):
19+
content = generator_func(*args, **kwargs)
20+
if isinstance(content, typing.AsyncIterable):
21+
content = _async_map(_wrap_event, content)
22+
else:
23+
content = map(_wrap_event, content)
24+
25+
return DatastarStreamingResponse(content)
26+
return _wrapper

0 commit comments

Comments
 (0)