DontPanicO / fastapi-distributed-websocket

A library to implement websocket for distibuted system based on FastAPI.
MIT License
54 stars 10 forks source link

Readme code does not reflect how methods should actually be used #30

Open SilentHacks opened 1 year ago

SilentHacks commented 1 year ago

Feature or enhancement

Make the README.md examples consistent with how the library actually works.

Issue

One of the examples of how to receive data in the README.md is as follows:

@app.websocket('/ws/{conn_id}')
async def websocket_endpoint(
    ws: WebSocket,
    conn_id: str,
    *,
    topic: Optional[Any] = None,
) -> None:
    connection: Connection = await manager.new_connection(ws, conn_id)
    try:
        while True:
            msg = await connection.receive_json()
            await manager.broadcast(msg)
    except WebSocketDisconnect:
        await manager.remove_connection(connection)

However, remove_connection() is not async so should not be awaited. Furthermore, the type of msg in the example is a dict and since the broadcast() method only accepts the Message type, this can't actually be used. The topic param isn't being passed to new_connection() either.

Similar story to the example after:

@app.websocket('/ws/{conn_id}')
async def websocket_endpoint(
    ws: WebSocket,
    conn_id: str,
    *,
    topic: Optional[Any] = None,
) -> None:
    connection: Connection = await manager.new_connection(ws, conn_id)
    # This is the preferred way of handling WebSocketDisconnect
    async for msg in connection.iter_json():
        await manager.receive(connection, msg)
    await manager.remove_connection(connection)

remove_connection() is not async. iter_json() is also returning dicts which don't fit with the receive() method. topic param is unused.

Pitch

The first example could become something like this:

@app.websocket('/ws/{conn_id}')
async def websocket_endpoint(
    ws: WebSocket,
    conn_id: str,
    *,
    topic: Optional[Any] = None,
) -> None:
    connection: Connection = await manager.new_connection(ws, conn_id, topic)
    try:
        while True:
            data = await connection.receive_json()
            msg = Message.from_client_message(data=data)
            manager.broadcast(msg)
    except WebSocketDisconnect:
        manager.remove_connection(connection)

There's a lot of potential solutions though. You could

For the second example, it could be something like this:

@app.websocket('/ws/{conn_id}')
async def websocket_endpoint(
    ws: WebSocket,
    conn_id: str,
    *,
    topic: Optional[Any] = None,
) -> None:
    connection: Connection = await manager.new_connection(ws, conn_id, topic)
    # This is the preferred way of handling WebSocketDisconnect
    async for msg in connection:
        await manager.receive(connection, msg)
    manager.remove_connection(connection)

Just a side-note but it could be cool to see some sort of context manager or putting into an iter so that it removes the connection on WebSocketDisconnect automatically.

Ryuk-me commented 1 year ago

https://github.com/DontPanicO/fastapi-distributed-websocket/issues/47 Can you please help? @SilentHacks