`Async for` can be used to iterate over a websocket's incoming messages

[bradlangel]

Is your feature request related to a problem? Please describe.

When creating a websocket server I'd like to use async for to iterate over the incoming messages from a connection. This is a feature that the websockets lib uses. Currently, if you try to use async for you get the following error:

TypeError: 'async for' requires an object with __aiter__ method, got WebsocketImplProtocol

Describe the solution you'd like

Ideally, I could run something like the following on a Sanic websocket route:

@app.websocket("/")
async def feed(request, ws):
    async for msg in ws:
        print(f'received: {msg.data}')
        await ws.send(msg.data)

Additional context This was originally discussed on the sanic-support channel on the discord server

[ChihweiLHBird]

I will take a look and see what can be done, if no one has started working on it. Thanks for creating the issue!

[hyansuper]

This feature was available in earlier version of sanic years ago, maybe earlier version uses websockets internally? The async for look very pythonic and I like it too. here's a simple workaround

class async_for_ws:
    def __init__(self, ws):
        self.ws = ws    
    def __aiter__(self):
        return self        
    async def __anext__(self):
        return await self.ws.recv()

async for msg in async_for_ws(ws):
    ....

[Tronic]

You can do it without __anext__ or any other extra boilerplate by simply adding this function to the websocket class:

async def __aiter__(self):
    while True:
        yield await self.recv()

As a user workaround, I suppose one can monkey patch that on Sanic's websocket, but hopefully @ChihweiLHBird gets the patch done and merged in time for the upcoming Sanic 22.6 release.

[ChihweiLHBird]

You can do it without __anext__ or any other extra boilerplate by simply adding this function to the websocket class:

async def __aiter__(self):
    while True:
        yield await self.recv()

As a user workaround, I suppose one can monkey patch that on Sanic's websocket, but hopefully @ChihweiLHBird gets the patch done and merged in time for the upcoming Sanic 22.6 release.

I think __anext__ is still worth because it makes the object supports anext. Consider this case:

async def test_anext():
    while True:
        await anext(ws)

[ChihweiLHBird]

My question is, should we make the whole ws (WebsocketImplProtocol) object async iterable? Or just make another recv-like method to return an async generator?

async def recv_generator(self):
    while True:
        yield await self.recv()

Any opinion?

[hyansuper]

I’d like it to have same behavior as the websockets lib

[ahopkins]

I’d like it to have same behavior as the websockets lib

:point_up:

[Tronic]

My question is, should we make the whole ws (WebsocketImplProtocol) object async iterable? Or just make another recv-like method to return an async generator?

The whole purpose of __aiter__ is to produce an iterable object. Making that function async lets Python do that built-in without needing us to implement an object for it, or supporting __anext__ on the websocket object. Similarly as you cannot next(a_list) without first doing iter(a_list) and then next() on what that returns. async for automatically calls first __aiter__ and then __anext__ behind the scenes.

[Tronic]

Because this is Sanic, I ran a benchmark. The built-in async iterator is some nanoseconds faster per call but while barely measurable, the difference is utterly negligible to anything else that receiving from a WebSocket does. The benchmark code (written on ipython prompt and timed with %timeit) is a bit too long to post here.

Also the situation reverses if only a few messages are to be received within the loop, as the built-in iterator still needs to be constructed, while returning self in __aiter__ avoids this.

One thing to consider is what to do if the websocket is closed, should it exit the loop cleanly or just raise some exception? Probably the latter, which doesn't then need any extra work. The loop will then never exit normally.

[ChihweiLHBird]

Because this is Sanic, I ran a benchmark. The built-in async iterator is some nanoseconds faster per call but while barely measurable, the difference is utterly negligible to anything else that receiving from a WebSocket does. The benchmark code (written on ipython prompt and timed with %timeit) is a bit too long to post here.

Also the situation reverses if only a few messages are to be received within the loop, as the built-in iterator still needs to be constructed, while returning self in __aiter__ avoids this.

One thing to consider is what to do if the websocket is closed, should it exit the loop cleanly or just raise some exception? Probably the latter, which doesn't then need any extra work. The loop will then never exit normally.

Okay, I agree with you. Benchmark matters and I think the scenario of large number of messages is more common.

ConnectionClosedOK exception will be raised if it closed normally. So, I think we can catch this and return while leaving other exceptions to raise up.

[Tronic]

Code clarity should probably take precedence since there really is no speed difference either way :)