aio-libs / aiohttp

Asynchronous HTTP client/server framework for asyncio and Python
https://docs.aiohttp.org
Other
14.97k stars 2k forks source link

ClientWebSocketResponse: Add way to get close message when using msg iterator #7682

Open garyvdm opened 11 months ago

garyvdm commented 11 months ago

Is your feature request related to a problem?

The websocket server I'm connecting to will send a close message. I want to get this message.

There doesn't seem to a way to do this when using the async message iterator on ClientWebSocketResponse:

e.g.

async with session.ws_connect('http://example.org/ws') as ws:
    async for msg in ws:
        if msg.type == aiohttp.WSMsgType.TEXT:
            await ws.send_str(msg.data + '/answer')
        elif msg.type == aiohttp.WSMsgType.ERROR:
            break
    # No way to get close message here

Describe the solution you'd like

There should be a ClientWebSocketResponse.close_message property that is set when the server has sent a close message

Describe alternatives you've considered

As a workaround, I'm using the ClientWebSocketResponse.receive() method instead of async message iterator. Then I'm able to handle the close message my self. E.g.:

close_msg = None
async with session.ws_connect('http://example.org/ws') as ws:
    while True:
        msg = await ws.receive()
        if msg.type == aiohttp.WSMsgType.TEXT:
            await ws.send_str(msg.data + '/answer')
        elif msg.type == aiohttp.WSMsgType.ERROR:
            break
        elif msg.type == WSMsgType.CLOSE:
            close_msg = msg.extra
            break
    logger.info(f"Closed: {ws.close_code} {close_msg}")

Related component

Client

Additional context

No response

Code of Conduct

Dreamsorcerer commented 2 weeks ago

I think a .close_reason attribute would be reasonable. Want to open a PR and work on this?