File size: 2,069 Bytes
5190f6d c70f9b8 5190f6d c70f9b8 5190f6d c70f9b8 5190f6d c70f9b8 5190f6d c70f9b8 5190f6d c70f9b8 5190f6d c70f9b8 5190f6d | 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 55 56 57 58 59 60 | from fastapi import WebSocket
from datetime import datetime
from typing import List
class Connection:
websocket: WebSocket
connection_time: datetime
def __init__(self, websocket: WebSocket, connection_time: datetime):
self.websocket = websocket
self.connection_time = connection_time
class ConnectionManager:
timeout = 60 * 5 # 5 minutes
def __init__(self):
self.active_connections: List[Connection] = []
async def connect(self, websocket: WebSocket):
print('Connecting')
await websocket.accept()
# Add connection time and websocket to active connections
self.active_connections.append(Connection(websocket=websocket, connection_time=datetime.now()))
def isConnected(self, websocket: WebSocket):
for connection in self.active_connections:
if connection.websocket == websocket:
return True
return False
def shouldDisconnect(self, websocket: WebSocket):
for connection in self.active_connections:
if connection.websocket == websocket:
if (datetime.now() - connection.connection_time).total_seconds() > self.timeout:
print('Disconnecting...')
return True
return False
async def receive_json(self, websocket: WebSocket):
if not self.isConnected(websocket):
return None
print('Receiving...')
data = await websocket.receive_json()
print('Received')
return data
def disconnect(self, websocket: WebSocket):
print('Disconnecting...')
for connection in self.active_connections:
if connection.websocket == websocket:
self.active_connections.remove(connection)
return True
return False
async def send_json(self, json, websocket: WebSocket):
print('Sending JSON...')
# Only send the message if the connection is still active
if self.isConnected(websocket):
await websocket.send_json(json) |