feat: Add simplistic websocket channels.

src/meta/sockets.py

@@ -0,0 +1,68 @@
+from abc import ABC
+from collections import defaultdict
+import json
+from typing import Any
+import logging
+
+import websockets
+
+
+logger = logging.getLogger(__name__)
+
+
+
+class Channel(ABC):
+    """
+    A channel is a stateful connection handler for a group of connected websockets.
+    """
+    name = "Root Channel"
+
+    def __init__(self, **kwargs):
+        self.connections = set()
+
+    @property
+    def empty(self):
+        return not self.connections
+
+    async def on_connection(self, websocket: websockets.WebSocketServerProtocol, event: dict[str, Any]):
+        logger.info(f"Channel '{self.name}' attached new connection {websocket=} {event=}")
+        self.connections.add(websocket)
+
+    async def del_connection(self, websocket: websockets.WebSocketServerProtocol):
+        logger.info(f"Channel '{self.name}' dropped connection {websocket=}")
+        self.connections.discard(websocket)
+
+    async def handle_message(self, websocket: websockets.WebSocketServerProtocol, message):
+        raise NotImplementedError
+
+    async def send_event(self, event, websocket=None):
+        message = json.dumps(event)
+        if not websocket:
+            for ws in self.connections:
+                await ws.send(message)
+        else:
+            await websocket.send(message)
+
+channels = {}
+
+def register_channel(name, channel: Channel):
+    channels[name] = channel
+
+
+async def root_handler(websocket: websockets.WebSocketServerProtocol):
+    message = await websocket.recv()
+    event = json.loads(message)
+
+    if event.get('type', None) != 'init':
+        raise ValueError("Received Websocket connection with no init.")
+
+    if (channel_name := event.get('channel', None)) not in channels:
+        raise ValueError(f"Received Init for unhandled channel {channel_name=}")
+    channel = channels[channel_name]
+
+    try:
+        await channel.on_connection(websocket, event)
+        async for message in websocket:
+            await channel.handle_message(websocket, message)
+    finally:
+        await channel.del_connection(websocket)