diff options
Diffstat (limited to 'pysite/websockets.py')
-rw-r--r-- | pysite/websockets.py | 123 |
1 files changed, 0 insertions, 123 deletions
diff --git a/pysite/websockets.py b/pysite/websockets.py deleted file mode 100644 index 213daace..00000000 --- a/pysite/websockets.py +++ /dev/null @@ -1,123 +0,0 @@ -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 |