from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi import APIRouter, Response
router = APIRouter()
html = """
Chat
WebSocket Chat
Your ID:
"""
class ConnectionManager:
def __init__(self) -> None:
self.active_connections: list[WebSocket] = []
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 send_personal_message(self, message: str, websocket: WebSocket) -> None:
await websocket.send_text(message)
async def broadcast(self, message: str) -> None:
for connection in self.active_connections:
await connection.send_text(message)
manager = ConnectionManager()
@router.get("/ws_example")
async def get() -> HTMLResponse:
return HTMLResponse(html)
@router.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: int) -> None:
await manager.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await manager.send_personal_message(f"You wrote: {data}", websocket)
await manager.broadcast(f"Client #{client_id} says: {data}")
except WebSocketDisconnect:
manager.disconnect(websocket)
await manager.broadcast(f"Client #{client_id} left the chat")