Dan Albert 6c4b8c81ee Update mypy.
Needed so mypy can recognize the new Python 3.12 generic syntax.
2023-12-01 16:14:54 -08:00

55 lines
1.6 KiB
Python

import asyncio
from asyncio import wait, Future
from fastapi import APIRouter, WebSocket
from fastapi.encoders import jsonable_encoder
from .eventstream import EventStream
from .models import GameUpdateEventsJs
from .. import GameContext
router: APIRouter = APIRouter()
class ConnectionManager:
def __init__(self) -> None:
self.active_connections: list[WebSocket] = []
async def shutdown(self) -> None:
futures: list[Future[None]] = []
for connection in self.active_connections:
futures.append(asyncio.create_task(connection.close()))
await wait(futures)
async def connect(self, websocket: WebSocket) -> None:
await websocket.accept()
self.active_connections.append(websocket)
def disconnect(self, websocket: WebSocket) -> None:
self.active_connections.remove(websocket)
async def broadcast(self, events: GameUpdateEventsJs) -> None:
futures = []
for connection in self.active_connections:
futures.append(
asyncio.create_task(connection.send_json(jsonable_encoder(events)))
)
await wait(futures)
manager = ConnectionManager()
@router.websocket("/eventstream")
async def event_stream(websocket: WebSocket) -> None:
await manager.connect(websocket)
while True:
if not (events := await EventStream.get()).empty:
if events.shutting_down:
await manager.shutdown()
return
await manager.broadcast(
GameUpdateEventsJs.from_events(events, GameContext.get())
)