-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
54 lines (44 loc) · 1.45 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
app = FastAPI()
app.mount(
"/static",
StaticFiles(directory = "static"),
name = "static"
)
templates = Jinja2Templates(
directory = "templates"
)
# dictionary to store connected WebSocket clients
connected_users = {}
app.add_middleware(
CORSMiddleware,
allow_origins = ["*"],
allow_methods = ["*"],
allow_headers = ["*"],
allow_credentials = True,
)
@app.get("/", response_class = HTMLResponse)
async def main(request: Request):
return templates.TemplateResponse(
name = "chatting.jinja",
context = {'request': request}
)
@app.websocket("/ws/{user_id}")
async def websocket_endpoint(user_id: str, websocket: WebSocket):
await websocket.accept()
# store the WebSocket connection in the dictionary
connected_users[user_id] = websocket
try:
while True:
data = await websocket.receive_text()
# send the received data to the other user
for user, user_ws in connected_users.items():
if user != user_id:
await user_ws.send_text(data)
except WebSocketDisconnect:
# if a user disconnects, remove them from the dictionary
del connected_users[user_id]