aboutsummaryrefslogtreecommitdiffstats
path: root/pysite/websockets.py
diff options
context:
space:
mode:
Diffstat (limited to 'pysite/websockets.py')
-rw-r--r--pysite/websockets.py123
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