Pythonchat_server.py
Copy
1
import asyncio
2
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
3
from collections import defaultdict
4
 
5
app = FastAPI()
6
 
7
class ConnectionManager:
8
    def __init__(self):
9
        self.rooms: dict[str, list[WebSocket]] = defaultdict(list)
10
 
11
    async def connect(self, websocket: WebSocket, room_id: str):
12
        await websocket.accept()
13
        self.rooms[room_id].append(websocket)
14
 
15
    def disconnect(self, websocket: WebSocket, room_id: str):
16
        if websocket in self.rooms[room_id]:
17
            self.rooms[room_id].remove(websocket)
18
 
19
    async def broadcast(self, message: str, room_id: str, exclude: WebSocket | None = None):
20
        dead = []
21
        for ws in self.rooms[room_id]:
22
            if ws == exclude:
23
                continue
24
            try:
25
                await ws.send_text(message)
26
            except Exception:
27
                dead.append(ws)
28
        for ws in dead:
29
            self.disconnect(ws, room_id)
30
 
31
manager = ConnectionManager()
32
 
33
@app.websocket("/ws/{room_id}/{user_id}")
34
async def websocket_endpoint(websocket: WebSocket, room_id: str, user_id: str):
35
    await manager.connect(websocket, room_id)
36
    await manager.broadcast(f"[{user_id}] joined the room", room_id, exclude=websocket)
37
    try:
38
        while True:
39
            text = await websocket.receive_text()
40
            await manager.broadcast(f"[{user_id}]: {text}", room_id)
41
    except WebSocketDisconnect:
42
        manager.disconnect(websocket, room_id)
43
        await manager.broadcast(f"[{user_id}] left the room", room_id)
44
 
45
@app.get("/rooms/{room_id}/count")
46
async def get_room_count(room_id: str):
47
    return {"room_id": room_id, "connections": len(manager.rooms.get(room_id, []))}
48