Bidirectional communiation with websockets in Quart

Question

I want to be able to use a WebSocket in Quart to receive any messages that are sent, and send any messages that I may need to send. There's no guarantee that messages will alternate between sending and receiving.

As an example, Quart's tutorial page on WebSockets features the following snippet:

@app.websocket('/api/v2/ws')
@collect_websocket
async def ws(queue):
    while True:
        data = await queue.get()
        await websocket.send(data)

Somehow, I want to modify the code in the while True loop so that I can check if there is any data to be received, but if there isn't, I will instead check the queue.

I would like to be able to await receiving on the socket only if there is something to receive (which maybe could be achieved if there were a timeout parameter in the receive method), but that is not an option.

So, how can I await a WebSocket for updates while also awaiting something else for updates?

Answer 1

The author of Quart answers this in the post Websockets in Quart, which contains a snippet which I've modified slightly to get the following:

import asyncio

from quart import copy_current_websocket_context, Quart, websocket

app = Quart(__name__)

@app.websocket('/ws')
async def ws():

    async def consumer():
        while True:
            data = await websocket.receive()

    async def producer():
        while True:
            await asyncio.sleep(1)
            await websocket.send(b'Message')

    consumer_task = asyncio.ensure_future(consumer())
    producer_task = asyncio.ensure_future(producer())
    try:
        await asyncio.gather(consumer_task, producer_task)
    finally:
        consumer_task.cancel()
        producer_task.cancel()

The snippet creates two different async functions with their own while True loops. Then, Python's asyncio.ensure_future is used to create two different Tasks to work on . Finally, asyncio.gather is called to evaluate the tasks concurrently.

By defining the two tasks inside the definition of ws, they act as closures, meaning they have access to the websocket "special" global object, which is only meaningful inside the ws function. If you wanted to define these functions outside the body of the ws function, perhaps because you need to call them from elsewhere, you could use the copy_current_websocket_context function from Quart when passing them to ensure_future:

consumer_task = asyncio.ensure_future(
    copy_current_websocket_context(consumer)()
)