53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
from __future__ import annotations
|
||
|
||
from fastapi import APIRouter, WebSocket, status
|
||
from fastapi.exceptions import WebSocketException
|
||
from starlette.websockets import WebSocketDisconnect
|
||
|
||
from app.ws.manager import ConnectionManager
|
||
|
||
router = APIRouter()
|
||
manager = ConnectionManager()
|
||
|
||
DEFAULT_HZ = 10
|
||
MIN_HZ = 1
|
||
MAX_HZ = 60
|
||
|
||
|
||
def _parse_hz(ws: WebSocket) -> int:
|
||
raw = ws.query_params.get("hz")
|
||
if raw is None:
|
||
return DEFAULT_HZ
|
||
try:
|
||
hz = int(raw)
|
||
except ValueError:
|
||
raise WebSocketException(
|
||
code=status.WS_1008_POLICY_VIOLATION, reason="Invalid 'hz' (int expected)"
|
||
)
|
||
if hz < MIN_HZ or hz > MAX_HZ:
|
||
raise WebSocketException(
|
||
code=status.WS_1008_POLICY_VIOLATION,
|
||
reason=f"Invalid 'hz' (allowed {MIN_HZ}..{MAX_HZ})",
|
||
)
|
||
return hz
|
||
|
||
|
||
@router.websocket("/ws/live")
|
||
async def ws_live(ws: WebSocket) -> None:
|
||
hz = _parse_hz(ws)
|
||
|
||
await manager.connect(ws, hz=hz)
|
||
try:
|
||
# Не обязательно принимать сообщения от клиента
|
||
# Но чтобы корректно ловить disconnect в некоторых клиентах - держим receive loop
|
||
while True:
|
||
await ws.receive_text()
|
||
except WebSocketDisconnect:
|
||
await manager.disconnect(ws)
|
||
except WebSocketException:
|
||
# если прилетит exception после accept — корректно удалим
|
||
await manager.disconnect(ws)
|
||
raise
|
||
except Exception:
|
||
await manager.disconnect(ws)
|