from falcon import WebSocketDisconnected
+from .static import TreeFileApp
-class BaseWebSocketApp:
- def __init__(self, hubapp):
- self.hubapp = hubapp
- self.conn = self.hubapp.app.hubapps["root"].conn
- def task_done(self):
- self.task = None
+class WebSocketApp:
+ def __init__(self, hubapp: TreeFileApp):
+ 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):
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}),
)
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", ())
]