| 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 |
|
|
| def __init__(self): |
| self.active_connections: List[Connection] = [] |
|
|
| async def connect(self, websocket: WebSocket): |
| print('Connecting') |
| await websocket.accept() |
| |
| 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...') |
| |
| if self.isConnected(websocket): |
| await websocket.send_json(json) |
