-
-
Notifications
You must be signed in to change notification settings - Fork 1.5k
/
connection.py
90 lines (71 loc) · 2.29 KB
/
connection.py
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
from typing import (
Any,
Awaitable,
Callable,
Dict,
List,
MutableMapping,
Optional,
Union,
)
from sanic.exceptions import InvalidUsage
ASGIMessage = MutableMapping[str, Any]
class WebSocketConnection:
"""
This is for ASGI Connections.
It provides an interface similar to WebsocketProtocol, but
sends/receives over an ASGI connection.
"""
# TODO
# - Implement ping/pong
def __init__(
self,
send: Callable[[ASGIMessage], Awaitable[None]],
receive: Callable[[], Awaitable[ASGIMessage]],
subprotocols: Optional[List[str]] = None,
) -> None:
self._send = send
self._receive = receive
self._subprotocols = subprotocols or []
async def send(self, data: Union[str, bytes], *args, **kwargs) -> None:
message: Dict[str, Union[str, bytes]] = {"type": "websocket.send"}
if isinstance(data, bytes):
message.update({"bytes": data})
else:
message.update({"text": str(data)})
await self._send(message)
async def recv(self, *args, **kwargs) -> Optional[str]:
message = await self._receive()
if message["type"] == "websocket.receive":
try:
return message["text"]
except KeyError:
try:
return message["bytes"].decode()
except KeyError:
raise InvalidUsage("Bad ASGI message received")
elif message["type"] == "websocket.disconnect":
pass
return None
receive = recv
async def accept(self, subprotocols: Optional[List[str]] = None) -> None:
subprotocol = None
if subprotocols:
for subp in subprotocols:
if subp in self.subprotocols:
subprotocol = subp
break
await self._send(
{
"type": "websocket.accept",
"subprotocol": subprotocol,
}
)
async def close(self, code: int = 1000, reason: str = "") -> None:
pass
@property
def subprotocols(self):
return self._subprotocols
@subprotocols.setter
def subprotocols(self, subprotocols: Optional[List[str]] = None):
self._subprotocols = subprotocols or []