diff --git a/pizzicore/pizzicore.py b/pizzicore/pizzicore.py index 30539b1..dc91cd9 100644 --- a/pizzicore/pizzicore.py +++ b/pizzicore/pizzicore.py @@ -1,4 +1,5 @@ import secrets +import logging import dbm from collections import defaultdict @@ -9,8 +10,8 @@ from fastapi import ( HTTPException, Depends, status, - WebSocketDisconnect, ) +from starlette.websockets import WebSocketDisconnect from fastapi.security import HTTPBasic, HTTPBasicCredentials from pydantic import BaseModel @@ -136,8 +137,14 @@ async def websocket_counter(websocket: WebSocket, cid: int): while True: try: - await websocket.send_text(str(counter_store.get(cid))) - except: + val = counter_store.get(cid) + await websocket.send_json(Value(counter=cid, value=val).json()) + except WebSocketDisconnect: + logging.debug("client disconnected") + manager.unsubscribe(cid, q) + return + except Exception: + logging.exception("unexpected error") manager.unsubscribe(cid, q) return await q.get()