From 5ed48ccdc8705f1dcffe7044b3980f9d5f7c95b9 Mon Sep 17 00:00:00 2001 From: Rupsi Kaushik Date: Sun, 9 Aug 2020 09:52:19 -0400 Subject: [PATCH] Export WebSocketDisconnect and add example handling disconnections to docs (#1822) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Sebastián Ramírez --- docs/en/docs/advanced/websockets.md | 27 ++++++ docs_src/websockets/tutorial003.py | 83 +++++++++++++++++++ fastapi/__init__.py | 2 +- .../test_websockets/test_tutorial003.py | 22 +++++ 4 files changed, 133 insertions(+), 1 deletion(-) create mode 100644 docs_src/websockets/tutorial003.py create mode 100644 tests/test_tutorial/test_websockets/test_tutorial003.py diff --git a/docs/en/docs/advanced/websockets.md b/docs/en/docs/advanced/websockets.md index 083c75752..1fe32bc61 100644 --- a/docs/en/docs/advanced/websockets.md +++ b/docs/en/docs/advanced/websockets.md @@ -137,6 +137,33 @@ With that you can connect the WebSocket and then send and receive messages: +## Handling disconnections and multiple clients + +When a WebSocket connection is closed, the `await websocket.receive_text()` will raise a `WebSocketDisconnect` exception, which you can then catch and handle like in this example. + +```Python hl_lines="81-83" +{!../../../docs_src/websockets/tutorial003.py!} +``` + +To try it out: + +* Open the app with several browser tabs. +* Write messages from them. +* Then close one of the tabs. + +That will raise the `WebSocketDisconnect` exception, and all the other clients will receive a message like: + +``` +Client #1596980209979 left the chat +``` + +!!! tip + The app above is a minimal and simple example to demonstrate how to handle and broadcast messages to several WebSocket connections. + + But have in mind that, as everything is handled in memory, in a single list, it will only work while the process is running, and will only work with a single process. + + If you need something easy to integrate with FastAPI but that is more robust, supported by Redis, PostgreSQL or others, check encode/broadcaster. + ## More info To learn more about the options, check Starlette's documentation for: diff --git a/docs_src/websockets/tutorial003.py b/docs_src/websockets/tutorial003.py new file mode 100644 index 000000000..d561633a8 --- /dev/null +++ b/docs_src/websockets/tutorial003.py @@ -0,0 +1,83 @@ +from typing import List + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from fastapi.responses import HTMLResponse + +app = FastAPI() + +html = """ + + + + Chat + + +

WebSocket Chat

+

Your ID:

+
+ + +
+ + + + +""" + + +class ConnectionManager: + def __init__(self): + self.active_connections: List[WebSocket] = [] + + async def connect(self, websocket: WebSocket): + await websocket.accept() + self.active_connections.append(websocket) + + def disconnect(self, websocket: WebSocket): + self.active_connections.remove(websocket) + + async def send_personal_message(self, message: str, websocket: WebSocket): + await websocket.send_text(message) + + async def broadcast(self, message: str): + for connection in self.active_connections: + await connection.send_text(message) + + +manager = ConnectionManager() + + +@app.get("/") +async def get(): + return HTMLResponse(html) + + +@app.websocket("/ws/{client_id}") +async def websocket_endpoint(websocket: WebSocket, client_id: int): + await manager.connect(websocket) + try: + while True: + data = await websocket.receive_text() + await manager.send_personal_message(f"You wrote: {data}", websocket) + await manager.broadcast(f"Client #{client_id} says: {data}") + except WebSocketDisconnect: + manager.disconnect(websocket) + await manager.broadcast(f"Client #{client_id} left the chat") diff --git a/fastapi/__init__.py b/fastapi/__init__.py index 5edd8a547..408c99db1 100644 --- a/fastapi/__init__.py +++ b/fastapi/__init__.py @@ -22,4 +22,4 @@ from .param_functions import ( from .requests import Request from .responses import Response from .routing import APIRouter -from .websockets import WebSocket +from .websockets import WebSocket, WebSocketDisconnect diff --git a/tests/test_tutorial/test_websockets/test_tutorial003.py b/tests/test_tutorial/test_websockets/test_tutorial003.py new file mode 100644 index 000000000..adc2cdae7 --- /dev/null +++ b/tests/test_tutorial/test_websockets/test_tutorial003.py @@ -0,0 +1,22 @@ +from fastapi.testclient import TestClient + +from docs_src.websockets.tutorial003 import app + +client = TestClient(app) + + +def test_websocket_handle_disconnection(): + with client.websocket_connect("/ws/1234") as connection, client.websocket_connect( + "/ws/5678" + ) as connection_two: + connection.send_text("Hello from 1234") + data1 = connection.receive_text() + assert data1 == "You wrote: Hello from 1234" + data2 = connection_two.receive_text() + client1_says = "Client #1234 says: Hello from 1234" + assert data2 == client1_says + data1 = connection.receive_text() + assert data1 == client1_says + connection_two.close() + data1 = connection.receive_text() + assert data1 == "Client #5678 left the chat"