X-Git-Url: https://git.mar77i.info/?a=blobdiff_plain;f=hub%2Fwebsocket.py;fp=hub%2Fwebsocket.py;h=89a7e2f3cfa4fd79baed82bc75c6581620eccab5;hb=a12839ce8fe46f0c2c0e98a37deebba05ea404b5;hp=0ecf9875e29a5ba5f66665a1b54ef30bd2a41f9a;hpb=3c5ec422ace644d848d2f845b0f3ef8de73462ef;p=hublib diff --git a/hub/websocket.py b/hub/websocket.py index 0ecf987..89a7e2f 100644 --- a/hub/websocket.py +++ b/hub/websocket.py @@ -9,24 +9,24 @@ from traceback import print_exception from falcon import WebSocketDisconnected -class BaseWebSocketApp: +class WebSocketApp: def __init__(self, hubapp): - self.hubapp = hubapp - self.conn = self.hubapp.app.hubapps["root"].conn - - def task_done(self): - self.task = None + self.name = hubapp.name + self.conn = hubapp.root.conn @staticmethod - async def process_websocket(conn, web_socket, extra_data={}, recipients=[]): + async def process_websocket(conn, web_socket, extra_data=None, recipients=None): try: while True: data = json.loads(await web_socket.receive_text()) - data.update(extra_data) + if extra_data: + data.update(extra_data) if callable(recipients): current_recipients = recipients(data) - else: + elif recipients: current_recipients = recipients + else: + raise ValueError("no recipients specified") for recipient in current_recipients: await conn.publish(recipient, pickle.dumps(data)) except (CancelledError, WebSocketDisconnected): @@ -76,11 +76,9 @@ class BaseWebSocketApp: if callable(leave_cb): await leave_cb() - -class WebSocketApp(BaseWebSocketApp): - async def join_leave_client_notify(self, redis, action, client_id): + async def client_notify(self, redis, action, client_id): await redis.publish( - f"{self.hubapp.name}-master", + f"{self.name}-master", pickle.dumps({"action": action, "client_id": client_id}), ) @@ -89,25 +87,25 @@ class WebSocketApp(BaseWebSocketApp): return await self.on_websocket( req, web_socket, - f"{self.hubapp.name}-client-{client_id}", + f"{self.name}-client-{client_id}", { "extra_data": {"client_id": client_id}, - "recipients": [f"{self.hubapp.name}-master"], + "recipients": [f"{self.name}-master"], }, - partial(self.join_leave_client_notify, self.conn, "join", client_id), - partial(self.join_leave_client_notify, self.conn, "leave", client_id), + partial(self.client_notify, self.conn, "join", client_id), + partial(self.client_notify, self.conn, "leave", client_id), ) async def on_websocket_master(self, req, web_socket): return await self.on_websocket( req, web_socket, - f"{self.hubapp.name}-master", + f"{self.name}-master", {"recipients": self.get_master_recipients}, ) def get_master_recipients(self, data): return [ - f"{self.hubapp.name}-client-{int(client_id)}" + f"{self.name}-client-{int(client_id)}" for client_id in data.pop("client_ids", ()) ]