2

I need help with the python web frame work, Quart, more specifically the websockets. I would like to be able to register a client when it connects (add it to a python list), and unregister them (remove it from the python list) when it disconnects. The closest thing I could find on the web is this code: 

connected = set()

async def handler(websocket, path):
    global connected
    # Register.
    connected.add(websocket)
    try:
        # Implement logic here.
        await asyncio.wait([ws.send("Hello!") for ws in connected])
        await asyncio.sleep(10)
    finally:
        # Unregister.
        connected.remove(websocket)

source

But this does not work with quart websockets.

Help would be appreciated.

Matthieu Brucher
  • 21,634
  • 7
  • 38
  • 62
Dup Dup
  • 55
  • 9

1 Answers1

6

This decorator when used to wrap a websocket handler, will add and remove websockets from the connected set. The _get_current_object method of the websocket is required to get the websocket in the current context, and the try-finally is required to ensure the websocket is removed regardless of any errors that are raised. Note the app.websocket must wrap (be before) the collect_websocket usage.

from functools import wraps

connected = set()

def collect_websocket(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        global connected
        connected.add(websocket._get_current_object())
        try:
            return await func(*args, **kwargs)
        finally:
            connected.remove(websocket._get_current_object())
    return wrapper                                                                                                                                                                                                            


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

Edit: I am the Quart author.

Jamie Lindsey
  • 928
  • 14
  • 26
pgjones
  • 6,044
  • 1
  • 14
  • 12
  • 1
    Hello, welcome to Stack Overflow! Please just note if you want to promote or recommend your own product/blog, there are some [guidelines in place](https://stackoverflow.com/help/promotion) for doing so. Following them will help you avoid giving the impression that you're spamming. Could you please [edit] to explicitly state your affiliation? Thanks. (If you're not actually affiliated, it may be worth mentioning that as well.) – DavidPostill Apr 15 '18 at 11:31
  • 1
    I've made an edit to state my involvement in Quart. I'm not sure what this answer would be promoting though? – pgjones Apr 15 '18 at 12:43
  • That's a fine edit. The promotion would be promoting your own code/blog ... – DavidPostill Apr 15 '18 at 13:00
  • I get this error: line 20, in collect_websocket @wraps(func) NameError: name 'wraps' is not defined, Is there something im missing? – Dup Dup Apr 15 '18 at 16:59
  • 1
    `from functools import wraps` probably, [see docs](https://docs.python.org/3/library/functools.html#functools.wraps) – pgjones Apr 15 '18 at 19:38
  • 1
    The OP was asking specifically about quart, so that edit was unnecessary. (It would make sense to add a disclaimer if the question didn't specify a framework, but it did). Thanks for your hard work @pgjones. – Clément May 31 '20 at 02:49