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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
|
import json
from flask import Blueprint
from geventwebsocket.websocket import WebSocket
class WS:
"""
Base class for representing a Websocket.
At minimum, you must implement the `on_message(self, message)` function. Without it, you won't be able to handle
any messages, and an error will be thrown!
If you need access to the database, you can mix-in DBMixin, just like any view class:
>>> class DBWebsocket(WS, DBMixin):
... name = "db_websocket"
... path = "/db_websocket" # This will be prefixed with "/ws" by the blueprint
... table = "ws"
...
... def on_message(self, message):
... self.send(
... json.loads(self.db.get(self.table_name, message))
... )
Please note that an instance of this class is created for every websocket connected to the path. This does, however,
mean that you can store any state required by your websocket.
"""
path = "" # type: str
name = "" # type: str
_connections = None
def __init__(self, socket: WebSocket):
self.socket = socket
def __new__(cls, *args, **kwargs):
if cls._connections is None:
cls._connections = []
return super().__new__(cls)
def on_open(self):
"""
Called once when the websocket is opened. Optional.
"""
def on_message(self, message: str):
"""
Called when a message is received by the websocket.
"""
raise NotImplementedError()
def on_close(self):
"""
Called once when the websocket is closed. Optional.
"""
def send(self, message, binary=None):
"""
Send a message to the currently-connected websocket, if it's open.
Nothing will happen if the websocket is closed.
"""
if not self.socket.closed:
self.socket.send(message, binary=binary)
def send_json(self, data):
return self.send(json.dumps(data))
@classmethod
def send_all(cls, message, binary=None):
for connection in cls._connections:
connection.send(message, binary=binary)
@classmethod
def send_all_json(cls, data):
for connection in cls._connections:
connection.send_json(data)
@classmethod
def setup(cls: "type(WS)", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint):
"""
Set up the websocket object, calling `setup()` on any superclasses as necessary (for example, on the DB
mixin).
This function will set up a websocket handler so that it behaves in a class-oriented way. It's up to you to
deal with message handling yourself, however.
"""
if hasattr(super(), "setup"):
super().setup(manager, blueprint)
if not cls.path or not cls.name:
raise RuntimeError("Websockets must have both `path` and `name` defined")
cls.manager = manager
def handle(socket: WebSocket):
"""
Wrap the current WS class, dispatching events to it as necessary. We're using gevent, so there's
no need to worry about blocking here.
"""
ws = cls(socket) # Instantiate the current class, passing it the WS object
cls._connections.append(ws)
try:
ws.on_open() # Call the "on_open" handler
while not socket.closed: # As long as the socket is open...
message = socket.receive() # Wait for a message
if not socket.closed: # If the socket didn't just close (there's always a None message on closing)
ws.on_message(message) # Call the "on_message" handler
ws.on_close() # The socket just closed, call the "on_close" handler
finally:
cls._connections.remove(ws)
blueprint.route(cls.path)(handle) # Register the handling function to the WS blueprint
|
