✨ Export WebSocketState in fastapi.websockets

[matiuszka]

Trivial improvement about exposing WebSocketState. It will eliminate the need of importing it from starlette directly.

Previously:

from fastapi.websockets import WebSocket
from starlette.websockets import WebSocketState

After:

from fastapi.websockets import WebSocket, WebSocketState

[Kludex]

How do you use WebSocketState?

[matiuszka]

So there is an endpoint that is sending some data, usually very rarely but still, and is receiving nothing. So while we are not sending any data and the WebSocket connection was closed by the client, the task is still running. So to add functionality that if the connection is closed task will be closed also we are monitoring the WebSocket status, by simple .recevie.

The endpoint more or less looks like that:

@router.websocket("/subscribe/some/data/changes")
async def subscribe_connectors_status(
    websocket: WebSocket,
) -> None:
    async def send_some_data() -> None:
        last_data = await get_data()
        await websocket.send_json({"data": last_data})
        while True:
            await asyncio.sleep(1)
            new_data = await get_data()
            if new_data != last_data:
                # This happens very rarely that's why I would like to close all of these unused tasks.
                last_data = new_data
                await websocket.send_json({"foo": last_data})

    async def watch_status() -> None:
        while websocket.client_state == WebSocketState.CONNECTED:
            # We are not expecting any data, but this will change to the client state if the connection will be closed.
            await websocket.receive()

        tg.cancel_scope.cancel()

    await websocket.accept()
    async with create_task_group() as tg:
        tg.start_soon(watch_status)
        tg.start_soon(send_some_data)

[codecov]

Codecov Report

Merging #4376 (36dadce) into master (0bb8920) will not change coverage.
The diff coverage is 100.00%.

@@            Coverage Diff            @@
##            master     #4376   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          535       535           
  Lines        13825     13826    +1     
=========================================
+ Hits         13825     13826    +1     

Impacted Files	Coverage Δ
fastapi/websockets.py	100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[github-actions]

📝 Docs preview for commit 36dadce at: https://630fd8992dbce202ce9f181a--fastapi.netlify.app

[tiangolo]

Thanks for the contribution @matiuszka! 🍰

And thanks @Kludex for the help! 🙇

Knowing how you are using it made me feel it makes sense to have this in, so I'll take it. 🚀

Thanks! ☕