1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
|
import io
import re
import json
from typing import Optional, Any, Union
class ServerSentEvent:
"""
Helper class to format data for Server-Sent Events (SSE).
"""
_LINE_SEP_EXPR = re.compile(r"\r\n|\r|\n")
DEFAULT_SEPARATOR = "\r\n"
def __init__(
self,
data: Optional[Any] = None,
*,
event: Optional[str] = None,
id: Optional[str] = None,
retry: Optional[int] = None,
comment: Optional[str] = None,
sep: Optional[str] = None,
) -> None:
self.data = data
self.event = event
self.id = id
self.retry = retry
self.comment = comment
self._sep = sep if sep is not None else self.DEFAULT_SEPARATOR
def encode(self) -> bytes:
buffer = io.StringIO()
if self.comment is not None:
for chunk in self._LINE_SEP_EXPR.split(str(self.comment)):
buffer.write(f": {chunk}{self._sep}")
if self.id is not None:
# Clean newlines in the event id
buffer.write("id: " + self._LINE_SEP_EXPR.sub("", self.id) + self._sep)
if self.event is not None:
# Clean newlines in the event name
buffer.write(
"event: " + self._LINE_SEP_EXPR.sub("", self.event) + self._sep
)
if self.data is not None:
# Break multi-line data into multiple data: lines
for chunk in self._LINE_SEP_EXPR.split(str(self.data)):
buffer.write(f"data: {chunk}{self._sep}")
if self.retry is not None:
if not isinstance(self.retry, int):
raise TypeError("retry argument must be int")
buffer.write(f"retry: {self.retry}{self._sep}")
buffer.write(self._sep)
return buffer.getvalue().encode("utf-8")
class JSONServerSentEvent(ServerSentEvent):
"""
Helper class to format JSON data for Server-Sent Events (SSE).
"""
def __init__(
self,
data: Optional[Any] = None,
*args,
**kwargs,
) -> None:
super().__init__(
json.dumps(
data,
ensure_ascii=False,
allow_nan=False,
indent=None,
separators=(",", ":"),
)
if data is not None
else None,
*args,
**kwargs,
)
def ensure_bytes(data: Union[bytes, dict, ServerSentEvent, Any], sep: str) -> bytes:
if isinstance(data, bytes):
return data
if isinstance(data, ServerSentEvent):
return data.encode()
if isinstance(data, dict):
data["sep"] = sep
return ServerSentEvent(**data).encode()
return ServerSentEvent(str(data), sep=sep).encode()
|
