Added WebSocket notify API for spools

spoolman/ws.py

@@ -0,0 +1,48 @@
+"""Websocket functionality."""
+
+import logging
+
+from fastapi import WebSocket
+
+logger = logging.getLogger(__name__)
+
+
+class WebsocketManager:
+    """Websocket manager."""
+
+    def __init__(self) -> None:
+        """Initialize."""
+        self.connections: dict[str, set[WebSocket]] = {}
+
+    def connect(self, pool: tuple[str, ...], websocket: WebSocket) -> None:
+        """Connect a websocket."""
+        pool_str = ",".join(pool)
+        if pool_str not in self.connections:
+            self.connections[pool_str] = set()
+        self.connections[pool_str].add(websocket)
+        logger.info(
+            "Client %s is now listening on pool %s",
+            websocket.client.host if websocket.client else "?",
+            pool_str,
+        )
+
+    def disconnect(self, pool: tuple[str, ...], websocket: WebSocket) -> None:
+        """Disconnect a websocket."""
+        pool_str = ",".join(pool)
+        if pool_str in self.connections:
+            self.connections[pool_str].remove(websocket)
+        logger.info(
+            "Client %s has stopped listening on pool %s",
+            websocket.client.host if websocket.client else "?",
+            pool_str,
+        )
+
+    async def send(self, pool: tuple[str, ...], message: str) -> None:
+        """Send a message to all websockets in a pool."""
+        pool_str = ",".join(pool)
+        if pool_str in self.connections:
+            for websocket in self.connections[pool_str]:
+                await websocket.send_text(message)
+
+
+websocket_manager = WebsocketManager()