Search code examples
pythonwebsocketquart

Python Quart websocket, send data between two clients


Using Quart I am trying to receive data from one client via a websocket, then have the Quart websocket server send it to a different client via websocket.

The two clients will be alone sharing the same url, other pairs of clients will have their own urls. This echo test works for both clients individually:

@copilot_ext.websocket('/ws/<unique_id>')
async def ws(unique_id):
    while True:
        data = await websocket.receive()
        await websocket.send(f"echo {data}") 

I have tried broadcasting using the example here https://pgjones.gitlab.io/quart/tutorials/websocket_tutorial.html#broadcasting although I can catch and print the different websockets, have not had much luck sending data from one client to the other :(

connected_websockets = set()

def collect_websocket(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        global connected_websockets
        send_channel, receive_channel = trio.open_memory_channel(2)
        connected_websockets.add(send_channel)
        try:
            return await func(send_channel, *args, **kwargs)
        finally:
            connected_websockets.remove(send_channel)
    return wrapper

@copilot_ext.websocket('/ws/<unique_id>')
@collect_websocket
async def ws(que, unique_id):
    while True:
        data = await websocket.receive()
        for send_channel in connected_websockets:
            await send_channel.send(f"message {data}")
            print(send_channel)

Just storing the websocket object and iterating through them doesn't work either

connected_websockets = set()

@copilot_ext.websocket('/ws/<unique_id>')
async def ws(unique_id):
    global connected_websockets
    while True:
        data = await websocket.receive()
        connected_websockets.add(websocket)
        for websockett in connected_websockets:
            await websockett.send(f"message {data}")
            print(type(websockett))

Solution

  • I think this snippet can form the basis of what you want to achieve. The idea is that rooms are a collection of queues keyed by the room id. Then each connected client has a queue in the room which any other clients put messages to. The send_task then runs in the background to send any messages to the client that are on its queue. I hope this makes sense,

    import asyncio
    from collections import defaultdict
    
    from quart import Quart, websocket
    
    app = Quart(__name__)
    
    websocket_rooms = defaultdict(set)
    
    async def send_task(ws, queue):
        while True:
            message = await queue.get()
            await ws.send(message)
    
    @app.websocket("/ws/<id>/")
    async def ws(id):
        global websocket_rooms
        queue = asyncio.Queue()
        websocket_rooms[id].add(queue)
        try:
            task = asyncio.ensure_future(send_task(websocket._get_current_object(), queue))
            while True:
                message = await websocket.receive()
                for other in websocket_rooms[id]:
                    if other is not queue:
                        await other.put(message)
        finally:
            task.cancel()
            await task
            websocket_rooms[id].remove(queue)