umap/ws.py
2024-01-29 13:08:17 +01:00

37 lines
1,013 B
Python

import asyncio
import websockets
from websockets.server import serve
import json
# Just relay all messages to other connected peers for now
CONNECTIONS = set()
async def join_and_listen(websocket):
print(f"Someone joined: {id(websocket)}")
CONNECTIONS.add(websocket)
try:
async for message in websocket:
# recompute the peers-list at the time of message-sending.
# doing so beforehand would miss new connections
peers = CONNECTIONS - {websocket}
print(message)
print(peers)
websockets.broadcast(peers, message)
finally:
CONNECTIONS.remove(websocket)
async def handler(websocket):
message = await websocket.recv()
event = json.loads(message)
# The first event should always be 'join'
assert event["type"] == "join"
await join_and_listen(websocket)
async def main():
async with serve(handler, "localhost", 8001):
await asyncio.Future() # run forever
asyncio.run(main())