1
+ from fastapi import APIRouter
2
+ from starlette .websockets import WebSocket , WebSocketDisconnect
3
+
4
+ router = APIRouter (prefix = "/chat" )
5
+
6
+
7
+ class ConnectionManager :
8
+ def __init__ (self ):
9
+ self .active_connections : list [WebSocket ] = []
10
+
11
+ async def connect (self , websocket : WebSocket ):
12
+ await websocket .accept ()
13
+ self .active_connections .append (websocket )
14
+
15
+ def disconnect (self , websocket : WebSocket ):
16
+ self .active_connections .remove (websocket )
17
+
18
+ async def send_personal_message (self , message : str , websocket : WebSocket ):
19
+ await websocket .send_text (message )
20
+
21
+ async def broadcast (self , message : str ):
22
+ for connection in self .active_connections :
23
+ await connection .send_text (message )
24
+
25
+
26
+ manager = ConnectionManager ()
27
+
28
+
29
+ @router .websocket ("/{client_id}" )
30
+ async def websocket_endpoint (websocket : WebSocket , client_id : int ):
31
+ await manager .connect (websocket )
32
+ try :
33
+ while True :
34
+ data = await websocket .receive_text ()
35
+ await manager .send_personal_message (f"You wrote: { data } " , websocket )
36
+ await manager .broadcast (f"Client #{ client_id } says: { data } " )
37
+ except WebSocketDisconnect :
38
+ manager .disconnect (websocket )
39
+ await manager .broadcast (f"Client #{ client_id } left the chat" )
0 commit comments