aboutsummaryrefslogtreecommitdiffstats
path: root/pysite
diff options
context:
space:
mode:
authorGravatar Gareth Coles <[email protected]>2018-05-29 17:53:08 +0100
committerGravatar Gareth Coles <[email protected]>2018-05-29 17:53:08 +0100
commit45919eeb10929c8f5af1e763555c5c9caf8cf40a (patch)
tree8d5917182d28ddc5b285192401ccc05349300e57 /pysite
parentMerge remote-tracking branch 'origin/master' (diff)
WS objects now keep track of their connections
Diffstat (limited to 'pysite')
-rw-r--r--pysite/views/ws/bot.py3
-rw-r--r--pysite/websockets.py38
2 files changed, 31 insertions, 10 deletions
diff --git a/pysite/views/ws/bot.py b/pysite/views/ws/bot.py
index 24af74df..816e7579 100644
--- a/pysite/views/ws/bot.py
+++ b/pysite/views/ws/bot.py
@@ -54,6 +54,3 @@ class BotWebsocket(WS, DBMixin):
def on_close(self):
self.log.debug("Bot | WS closed.")
self.do_changefeed = False
-
- def send_json(self, data):
- return self.send(json.dumps(data))
diff --git a/pysite/websockets.py b/pysite/websockets.py
index d72f44b7..ff8f0a12 100644
--- a/pysite/websockets.py
+++ b/pysite/websockets.py
@@ -1,3 +1,5 @@
+import json
+
from flask import Blueprint
from geventwebsocket.websocket import WebSocket
@@ -28,9 +30,17 @@ class WS:
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.
@@ -58,6 +68,17 @@ class WS:
if not self.socket.closed:
self.socket.send(message, binary=binary)
+ def send_json(self, data):
+ return self.send(json.dumps(data))
+
+ def send_all(self, message, binary=None):
+ for connection in self._connections:
+ connection.send(message, binary=binary)
+
+ def send_all_json(self, data):
+ for connection in self._connections:
+ connection.send_json(data)
+
@classmethod
def setup(cls: "type(WS)", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint):
"""
@@ -83,15 +104,18 @@ class WS:
"""
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
- 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
+ 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
+ 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
+ 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