WebSockets

您可以在 FastAPI 中使用 WebSockets。

安装 WebSockets

首先，您需要安装 WebSockets：

  1. $ pip install websockets
  3. ---> 100%

WebSockets 客户端

在生产环境中

在您的生产系统中，您可能使用现代框架（如React、Vue.js或Angular）创建了一个前端。

要使用 WebSockets 与后端进行通信，您可能会使用前端的工具。

或者，您可能有一个原生移动应用程序，直接使用原生代码与 WebSocket 后端通信。

或者，您可能有其他与 WebSocket 终端通信的方式。

但是，在本示例中，我们将使用一个非常简单的HTML文档，其中包含一些JavaScript，全部放在一个长字符串中。

当然，这并不是最优的做法，您不应该在生产环境中使用它。

在生产环境中，您应该选择上述任一选项。

但这是一种专注于 WebSockets 的服务器端并提供一个工作示例的最简单方式：

  1. from fastapi import FastAPI, WebSocket
  2. from fastapi.responses import HTMLResponse
  4. app = FastAPI()
  6. html = """
  7. <!DOCTYPE html>
  8. <html>
  9.     <head>
  10.         <title>Chat</title>
  11.     </head>
  12.     <body>
  13.         <h1>WebSocket Chat</h1>
  14.         <form action="" onsubmit="sendMessage(event)">
  15.             <input type="text" id="messageText" autocomplete="off"/>
  16.             <button>Send</button>
  17.         </form>
  18.         <ul id='messages'>
  19.         </ul>
  20.         <script>
  21.             var ws = new WebSocket("ws://localhost:8000/ws");
  22.             ws.onmessage = function(event) {
  23.                 var messages = document.getElementById('messages')
  24.                 var message = document.createElement('li')
  25.                 var content = document.createTextNode(event.data)
  26.                 message.appendChild(content)
  27.                 messages.appendChild(message)
  28.             };
  29.             function sendMessage(event) {
  30.                 var input = document.getElementById("messageText")
  31.                 ws.send(input.value)
  32.                 input.value = ''
  33.                 event.preventDefault()
  34.             }
  35.         </script>
  36.     </body>
  37. </html>
  38. """
  41. @app.get("/")
  42. async def get():
  43.     return HTMLResponse(html)
  46. @app.websocket("/ws")
  47. async def websocket_endpoint(websocket: WebSocket):
  48.     await websocket.accept()
  49.     while True:
  50.         data = await websocket.receive_text()
  51.         await websocket.send_text(f"Message text was: {data}")

创建 websocket

在您的 FastAPI 应用程序中，创建一个 websocket：

  1. from fastapi import FastAPI, WebSocket
  2. from fastapi.responses import HTMLResponse
  4. app = FastAPI()
  6. html = """
  7. <!DOCTYPE html>
  8. <html>
  9.     <head>
  10.         <title>Chat</title>
  11.     </head>
  12.     <body>
  13.         <h1>WebSocket Chat</h1>
  14.         <form action="" onsubmit="sendMessage(event)">
  15.             <input type="text" id="messageText" autocomplete="off"/>
  16.             <button>Send</button>
  17.         </form>
  18.         <ul id='messages'>
  19.         </ul>
  20.         <script>
  21.             var ws = new WebSocket("ws://localhost:8000/ws");
  22.             ws.onmessage = function(event) {
  23.                 var messages = document.getElementById('messages')
  24.                 var message = document.createElement('li')
  25.                 var content = document.createTextNode(event.data)
  26.                 message.appendChild(content)
  27.                 messages.appendChild(message)
  28.             };
  29.             function sendMessage(event) {
  30.                 var input = document.getElementById("messageText")
  31.                 ws.send(input.value)
  32.                 input.value = ''
  33.                 event.preventDefault()
  34.             }
  35.         </script>
  36.     </body>
  37. </html>
  38. """
  41. @app.get("/")
  42. async def get():
  43.     return HTMLResponse(html)
  46. @app.websocket("/ws")
  47. async def websocket_endpoint(websocket: WebSocket):
  48.     await websocket.accept()
  49.     while True:
  50.         data = await websocket.receive_text()
  51.         await websocket.send_text(f"Message text was: {data}")

技术细节

您也可以使用 from starlette.websockets import WebSocket。

FastAPI 直接提供了相同的 WebSocket，只是为了方便开发人员。但它直接来自 Starlette。

等待消息并发送消息

在您的 WebSocket 路由中，您可以使用 await 等待消息并发送消息。

  1. from fastapi import FastAPI, WebSocket
  2. from fastapi.responses import HTMLResponse
  4. app = FastAPI()
  6. html = """
  7. <!DOCTYPE html>
  8. <html>
  9.     <head>
  10.         <title>Chat</title>
  11.     </head>
  12.     <body>
  13.         <h1>WebSocket Chat</h1>
  14.         <form action="" onsubmit="sendMessage(event)">
  15.             <input type="text" id="messageText" autocomplete="off"/>
  16.             <button>Send</button>
  17.         </form>
  18.         <ul id='messages'>
  19.         </ul>
  20.         <script>
  21.             var ws = new WebSocket("ws://localhost:8000/ws");
  22.             ws.onmessage = function(event) {
  23.                 var messages = document.getElementById('messages')
  24.                 var message = document.createElement('li')
  25.                 var content = document.createTextNode(event.data)
  26.                 message.appendChild(content)
  27.                 messages.appendChild(message)
  28.             };
  29.             function sendMessage(event) {
  30.                 var input = document.getElementById("messageText")
  31.                 ws.send(input.value)
  32.                 input.value = ''
  33.                 event.preventDefault()
  34.             }
  35.         </script>
  36.     </body>
  37. </html>
  38. """
  41. @app.get("/")
  42. async def get():
  43.     return HTMLResponse(html)
  46. @app.websocket("/ws")
  47. async def websocket_endpoint(websocket: WebSocket):
  48.     await websocket.accept()
  49.     while True:
  50.         data = await websocket.receive_text()
  51.         await websocket.send_text(f"Message text was: {data}")

您可以接收和发送二进制、文本和 JSON 数据。

尝试一下

如果您的文件名为 main.py，请使用以下命令运行应用程序：

  1. $ uvicorn main:app --reload
  3. <span style="color: green;">INFO</span>:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)

在浏览器中打开 http://127.0.0.1:8000。

您将看到一个简单的页面，如下所示：

WebSockets - 图1

您可以在输入框中输入消息并发送：

WebSockets - 图2

您的 FastAPI 应用程序将回复：

WebSockets - 图3

您可以发送（和接收）多条消息：

WebSockets - 图4

所有这些消息都将使用同一个 WebSocket 连

接。

使用 Depends 和其他依赖项

在 WebSocket 端点中，您可以从 fastapi 导入并使用以下内容：

  • Depends
  • Security
  • Cookie
  • Header
  • Path
  • Query

它们的工作方式与其他 FastAPI 端点/ 路径操作 相同：

Python 3.10+Python 3.9+Python 3.8+Python 3.10+ 非带注解版本Python 3.8+ 非带注解版本

  1. from typing import Annotated
  3. from fastapi import (
  4.     Cookie,
  5.     Depends,
  6.     FastAPI,
  7.     Query,
  8.     WebSocket,
  9.     WebSocketException,
  10.     status,
  11. )
  12. from fastapi.responses import HTMLResponse
  14. app = FastAPI()
  16. html = """
  17. <!DOCTYPE html>
  18. <html>
  19.     <head>
  20.         <title>Chat</title>
  21.     </head>
  22.     <body>
  23.         <h1>WebSocket Chat</h1>
  24.         <form action="" onsubmit="sendMessage(event)">
  25.             <label>Item ID: <input type="text" id="itemId" autocomplete="off" value="foo"/></label>
  26.             <label>Token: <input type="text" id="token" autocomplete="off" value="some-key-token"/></label>
  27.             <button onclick="connect(event)">Connect</button>
  28.             <hr>
  29.             <label>Message: <input type="text" id="messageText" autocomplete="off"/></label>
  30.             <button>Send</button>
  31.         </form>
  32.         <ul id='messages'>
  33.         </ul>
  34.         <script>
  35.         var ws = null;
  36.             function connect(event) {
  37.                 var itemId = document.getElementById("itemId")
  38.                 var token = document.getElementById("token")
  39.                 ws = new WebSocket("ws://localhost:8000/items/" + itemId.value + "/ws?token=" + token.value);
  40.                 ws.onmessage = function(event) {
  41.                     var messages = document.getElementById('messages')
  42.                     var message = document.createElement('li')
  43.                     var content = document.createTextNode(event.data)
  44.                     message.appendChild(content)
  45.                     messages.appendChild(message)
  46.                 };
  47.                 event.preventDefault()
  48.             }
  49.             function sendMessage(event) {
  50.                 var input = document.getElementById("messageText")
  51.                 ws.send(input.value)
  52.                 input.value = ''
  53.                 event.preventDefault()
  54.             }
  55.         </script>
  56.     </body>
  57. </html>
  58. """
  61. @app.get("/")
  62. async def get():
  63.     return HTMLResponse(html)
  66. async def get_cookie_or_token(
  67.     websocket: WebSocket,
  68.     session: Annotated[str | None, Cookie()] = None,
  69.     token: Annotated[str | None, Query()] = None,
  70. ):
  71.     if session is None and token is None:
  72.         raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
  73.     return session or token
  76. @app.websocket("/items/{item_id}/ws")
  77. async def websocket_endpoint(
  78.     *,
  79.     websocket: WebSocket,
  80.     item_id: str,
  81.     q: int | None = None,
  82.     cookie_or_token: Annotated[str, Depends(get_cookie_or_token)],
  83. ):
  84.     await websocket.accept()
  85.     while True:
  86.         data = await websocket.receive_text()
  87.         await websocket.send_text(
  88.             f"Session cookie or query token value is: {cookie_or_token}"
  89.         )
  90.         if q is not None:
  91.             await websocket.send_text(f"Query parameter q is: {q}")
  92.         await websocket.send_text(f"Message text was: {data}, for item ID: {item_id}")
  1. from typing import Annotated, Union
  3. from fastapi import (
  4.     Cookie,
  5.     Depends,
  6.     FastAPI,
  7.     Query,
  8.     WebSocket,
  9.     WebSocketException,
  10.     status,
  11. )
  12. from fastapi.responses import HTMLResponse
  14. app = FastAPI()
  16. html = """
  17. <!DOCTYPE html>
  18. <html>
  19.     <head>
  20.         <title>Chat</title>
  21.     </head>
  22.     <body>
  23.         <h1>WebSocket Chat</h1>
  24.         <form action="" onsubmit="sendMessage(event)">
  25.             <label>Item ID: <input type="text" id="itemId" autocomplete="off" value="foo"/></label>
  26.             <label>Token: <input type="text" id="token" autocomplete="off" value="some-key-token"/></label>
  27.             <button onclick="connect(event)">Connect</button>
  28.             <hr>
  29.             <label>Message: <input type="text" id="messageText" autocomplete="off"/></label>
  30.             <button>Send</button>
  31.         </form>
  32.         <ul id='messages'>
  33.         </ul>
  34.         <script>
  35.         var ws = null;
  36.             function connect(event) {
  37.                 var itemId = document.getElementById("itemId")
  38.                 var token = document.getElementById("token")
  39.                 ws = new WebSocket("ws://localhost:8000/items/" + itemId.value + "/ws?token=" + token.value);
  40.                 ws.onmessage = function(event) {
  41.                     var messages = document.getElementById('messages')
  42.                     var message = document.createElement('li')
  43.                     var content = document.createTextNode(event.data)
  44.                     message.appendChild(content)
  45.                     messages.appendChild(message)
  46.                 };
  47.                 event.preventDefault()
  48.             }
  49.             function sendMessage(event) {
  50.                 var input = document.getElementById("messageText")
  51.                 ws.send(input.value)
  52.                 input.value = ''
  53.                 event.preventDefault()
  54.             }
  55.         </script>
  56.     </body>
  57. </html>
  58. """
  61. @app.get("/")
  62. async def get():
  63.     return HTMLResponse(html)
  66. async def get_cookie_or_token(
  67.     websocket: WebSocket,
  68.     session: Annotated[Union[str, None], Cookie()] = None,
  69.     token: Annotated[Union[str, None], Query()] = None,
  70. ):
  71.     if session is None and token is None:
  72.         raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
  73.     return session or token
  76. @app.websocket("/items/{item_id}/ws")
  77. async def websocket_endpoint(
  78.     *,
  79.     websocket: WebSocket,
  80.     item_id: str,
  81.     q: Union[int, None] = None,
  82.     cookie_or_token: Annotated[str, Depends(get_cookie_or_token)],
  83. ):
  84.     await websocket.accept()
  85.     while True:
  86.         data = await websocket.receive_text()
  87.         await websocket.send_text(
  88.             f"Session cookie or query token value is: {cookie_or_token}"
  89.         )
  90.         if q is not None:
  91.             await websocket.send_text(f"Query parameter q is: {q}")
  92.         await websocket.send_text(f"Message text was: {data}, for item ID: {item_id}")
  1. from typing import Union
  3. from fastapi import (
  4.     Cookie,
  5.     Depends,
  6.     FastAPI,
  7.     Query,
  8.     WebSocket,
  9.     WebSocketException,
  10.     status,
  11. )
  12. from fastapi.responses import HTMLResponse
  13. from typing_extensions import Annotated
  15. app = FastAPI()
  17. html = """
  18. <!DOCTYPE html>
  19. <html>
  20.     <head>
  21.         <title>Chat</title>
  22.     </head>
  23.     <body>
  24.         <h1>WebSocket Chat</h1>
  25.         <form action="" onsubmit="sendMessage(event)">
  26.             <label>Item ID: <input type="text" id="itemId" autocomplete="off" value="foo"/></label>
  27.             <label>Token: <input type="text" id="token" autocomplete="off" value="some-key-token"/></label>
  28.             <button onclick="connect(event)">Connect</button>
  29.             <hr>
  30.             <label>Message: <input type="text" id="messageText" autocomplete="off"/></label>
  31.             <button>Send</button>
  32.         </form>
  33.         <ul id='messages'>
  34.         </ul>
  35.         <script>
  36.         var ws = null;
  37.             function connect(event) {
  38.                 var itemId = document.getElementById("itemId")
  39.                 var token = document.getElementById("token")
  40.                 ws = new WebSocket("ws://localhost:8000/items/" + itemId.value + "/ws?token=" + token.value);
  41.                 ws.onmessage = function(event) {
  42.                     var messages = document.getElementById('messages')
  43.                     var message = document.createElement('li')
  44.                     var content = document.createTextNode(event.data)
  45.                     message.appendChild(content)
  46.                     messages.appendChild(message)
  47.                 };
  48.                 event.preventDefault()
  49.             }
  50.             function sendMessage(event) {
  51.                 var input = document.getElementById("messageText")
  52.                 ws.send(input.value)
  53.                 input.value = ''
  54.                 event.preventDefault()
  55.             }
  56.         </script>
  57.     </body>
  58. </html>
  59. """
  62. @app.get("/")
  63. async def get():
  64.     return HTMLResponse(html)
  67. async def get_cookie_or_token(
  68.     websocket: WebSocket,
  69.     session: Annotated[Union[str, None], Cookie()] = None,
  70.     token: Annotated[Union[str, None], Query()] = None,
  71. ):
  72.     if session is None and token is None:
  73.         raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
  74.     return session or token
  77. @app.websocket("/items/{item_id}/ws")
  78. async def websocket_endpoint(
  79.     *,
  80.     websocket: WebSocket,
  81.     item_id: str,
  82.     q: Union[int, None] = None,
  83.     cookie_or_token: Annotated[str, Depends(get_cookie_or_token)],
  84. ):
  85.     await websocket.accept()
  86.     while True:
  87.         data = await websocket.receive_text()
  88.         await websocket.send_text(
  89.             f"Session cookie or query token value is: {cookie_or_token}"
  90.         )
  91.         if q is not None:
  92.             await websocket.send_text(f"Query parameter q is: {q}")
  93.         await websocket.send_text(f"Message text was: {data}, for item ID: {item_id}")

Tip

如果可能，请尽量使用 Annotated 版本。

  1. from fastapi import (
  2.     Cookie,
  3.     Depends,
  4.     FastAPI,
  5.     Query,
  6.     WebSocket,
  7.     WebSocketException,
  8.     status,
  9. )
  10. from fastapi.responses import HTMLResponse
  12. app = FastAPI()
  14. html = """
  15. <!DOCTYPE html>
  16. <html>
  17.     <head>
  18.         <title>Chat</title>
  19.     </head>
  20.     <body>
  21.         <h1>WebSocket Chat</h1>
  22.         <form action="" onsubmit="sendMessage(event)">
  23.             <label>Item ID: <input type="text" id="itemId" autocomplete="off" value="foo"/></label>
  24.             <label>Token: <input type="text" id="token" autocomplete="off" value="some-key-token"/></label>
  25.             <button onclick="connect(event)">Connect</button>
  26.             <hr>
  27.             <label>Message: <input type="text" id="messageText" autocomplete="off"/></label>
  28.             <button>Send</button>
  29.         </form>
  30.         <ul id='messages'>
  31.         </ul>
  32.         <script>
  33.         var ws = null;
  34.             function connect(event) {
  35.                 var itemId = document.getElementById("itemId")
  36.                 var token = document.getElementById("token")
  37.                 ws = new WebSocket("ws://localhost:8000/items/" + itemId.value + "/ws?token=" + token.value);
  38.                 ws.onmessage = function(event) {
  39.                     var messages = document.getElementById('messages')
  40.                     var message = document.createElement('li')
  41.                     var content = document.createTextNode(event.data)
  42.                     message.appendChild(content)
  43.                     messages.appendChild(message)
  44.                 };
  45.                 event.preventDefault()
  46.             }
  47.             function sendMessage(event) {
  48.                 var input = document.getElementById("messageText")
  49.                 ws.send(input.value)
  50.                 input.value = ''
  51.                 event.preventDefault()
  52.             }
  53.         </script>
  54.     </body>
  55. </html>
  56. """
  59. @app.get("/")
  60. async def get():
  61.     return HTMLResponse(html)
  64. async def get_cookie_or_token(
  65.     websocket: WebSocket,
  66.     session: str | None = Cookie(default=None),
  67.     token: str | None = Query(default=None),
  68. ):
  69.     if session is None and token is None:
  70.         raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
  71.     return session or token
  74. @app.websocket("/items/{item_id}/ws")
  75. async def websocket_endpoint(
  76.     websocket: WebSocket,
  77.     item_id: str,
  78.     q: int | None = None,
  79.     cookie_or_token: str = Depends(get_cookie_or_token),
  80. ):
  81.     await websocket.accept()
  82.     while True:
  83.         data = await websocket.receive_text()
  84.         await websocket.send_text(
  85.             f"Session cookie or query token value is: {cookie_or_token}"
  86.         )
  87.         if q is not None:
  88.             await websocket.send_text(f"Query parameter q is: {q}")
  89.         await websocket.send_text(f"Message text was: {data}, for item ID: {item_id}")

Tip

如果可能，请尽量使用 Annotated 版本。

  1. from typing import Union
  3. from fastapi import (
  4.     Cookie,
  5.     Depends,
  6.     FastAPI,
  7.     Query,
  8.     WebSocket,
  9.     WebSocketException,
  10.     status,
  11. )
  12. from fastapi.responses import HTMLResponse
  14. app = FastAPI()
  16. html = """
  17. <!DOCTYPE html>
  18. <html>
  19.     <head>
  20.         <title>Chat</title>
  21.     </head>
  22.     <body>
  23.         <h1>WebSocket Chat</h1>
  24.         <form action="" onsubmit="sendMessage(event)">
  25.             <label>Item ID: <input type="text" id="itemId" autocomplete="off" value="foo"/></label>
  26.             <label>Token: <input type="text" id="token" autocomplete="off" value="some-key-token"/></label>
  27.             <button onclick="connect(event)">Connect</button>
  28.             <hr>
  29.             <label>Message: <input type="text" id="messageText" autocomplete="off"/></label>
  30.             <button>Send</button>
  31.         </form>
  32.         <ul id='messages'>
  33.         </ul>
  34.         <script>
  35.         var ws = null;
  36.             function connect(event) {
  37.                 var itemId = document.getElementById("itemId")
  38.                 var token = document.getElementById("token")
  39.                 ws = new WebSocket("ws://localhost:8000/items/" + itemId.value + "/ws?token=" + token.value);
  40.                 ws.onmessage = function(event) {
  41.                     var messages = document.getElementById('messages')
  42.                     var message = document.createElement('li')
  43.                     var content = document.createTextNode(event.data)
  44.                     message.appendChild(content)
  45.                     messages.appendChild(message)
  46.                 };
  47.                 event.preventDefault()
  48.             }
  49.             function sendMessage(event) {
  50.                 var input = document.getElementById("messageText")
  51.                 ws.send(input.value)
  52.                 input.value = ''
  53.                 event.preventDefault()
  54.             }
  55.         </script>
  56.     </body>
  57. </html>
  58. """
  61. @app.get("/")
  62. async def get():
  63.     return HTMLResponse(html)
  66. async def get_cookie_or_token(
  67.     websocket: WebSocket,
  68.     session: Union[str, None] = Cookie(default=None),
  69.     token: Union[str, None] = Query(default=None),
  70. ):
  71.     if session is None and token is None:
  72.         raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
  73.     return session or token
  76. @app.websocket("/items/{item_id}/ws")
  77. async def websocket_endpoint(
  78.     websocket: WebSocket,
  79.     item_id: str,
  80.     q: Union[int, None] = None,
  81.     cookie_or_token: str = Depends(get_cookie_or_token),
  82. ):
  83.     await websocket.accept()
  84.     while True:
  85.         data = await websocket.receive_text()
  86.         await websocket.send_text(
  87.             f"Session cookie or query token value is: {cookie_or_token}"
  88.         )
  89.         if q is not None:
  90.             await websocket.send_text(f"Query parameter q is: {q}")
  91.         await websocket.send_text(f"Message text was: {data}, for item ID: {item_id}")

Info

由于这是一个 WebSocket，抛出 HTTPException 并不是很合理，而是抛出 WebSocketException。

您可以使用规范中定义的有效代码。

尝试带有依赖项的 WebSockets

如果您的文件名为 main.py，请使用以下命令运行应用程序：

  1. $ uvicorn main:app --reload
  3. <span style="color: green;">INFO</span>:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)

在浏览器中打开 http://127.0.0.1:8000。

在页面中，您可以设置：

  • “Item ID”，用于路径。
  • “Token”，作为查询参数。

Tip

注意，查询参数 token 将由依赖项处理。

通过这样，您可以连接 WebSocket，然后发送和接收消息：

WebSockets - 图5

处理断开连接和多个客户端

当 WebSocket 连接关闭时，await websocket.receive_text() 将引发 WebSocketDisconnect 异常，您可以捕获并处理该异常，就像本示例中的示例一样。

Python 3.9+Python 3.8+

  1. from fastapi import FastAPI, WebSocket, WebSocketDisconnect
  2. from fastapi.responses import HTMLResponse
  4. app = FastAPI()
  6. html = """
  7. <!DOCTYPE html>
  8. <html>
  9.     <head>
  10.         <title>Chat</title>
  11.     </head>
  12.     <body>
  13.         <h1>WebSocket Chat</h1>
  14.         <h2>Your ID: <span id="ws-id"></span></h2>
  15.         <form action="" onsubmit="sendMessage(event)">
  16.             <input type="text" id="messageText" autocomplete="off"/>
  17.             <button>Send</button>
  18.         </form>
  19.         <ul id='messages'>
  20.         </ul>
  21.         <script>
  22.             var client_id = Date.now()
  23.             document.querySelector("#ws-id").textContent = client_id;
  24.             var ws = new WebSocket(`ws://localhost:8000/ws/${client_id}`);
  25.             ws.onmessage = function(event) {
  26.                 var messages = document.getElementById('messages')
  27.                 var message = document.createElement('li')
  28.                 var content = document.createTextNode(event.data)
  29.                 message.appendChild(content)
  30.                 messages.appendChild(message)
  31.             };
  32.             function sendMessage(event) {
  33.                 var input = document.getElementById("messageText")
  34.                 ws.send(input.value)
  35.                 input.value = ''
  36.                 event.preventDefault()
  37.             }
  38.         </script>
  39.     </body>
  40. </html>
  41. """
  44. class ConnectionManager:
  45.     def __init__(self):
  46.         self.active_connections: list[WebSocket] = []
  48.     async def connect(self, websocket: WebSocket):
  49.         await websocket.accept()
  50.         self.active_connections.append(websocket)
  52.     def disconnect(self, websocket: WebSocket):
  53.         self.active_connections.remove(websocket)
  55.     async def send_personal_message(self, message: str, websocket: WebSocket):
  56.         await websocket.send_text(message)
  58.     async def broadcast(self, message: str):
  59.         for connection in self.active_connections:
  60.             await connection.send_text(message)
  63. manager = ConnectionManager()
  66. @app.get("/")
  67. async def get():
  68.     return HTMLResponse(html)
  71. @app.websocket("/ws/{client_id}")
  72. async def websocket_endpoint(websocket: WebSocket, client_id: int):
  73.     await manager.connect(websocket)
  74.     try:
  75.         while True:
  76.             data = await websocket.receive_text()
  77.             await manager.send_personal_message(f"You wrote: {data}", websocket)
  78.             await manager.broadcast(f"Client #{client_id} says: {data}")
  79.     except WebSocketDisconnect:
  80.         manager.disconnect(websocket)
  81.         await manager.broadcast(f"Client #{client_id} left the chat")
  1. from typing import List
  3. from fastapi import FastAPI, WebSocket, WebSocketDisconnect
  4. from fastapi.responses import HTMLResponse
  6. app = FastAPI()
  8. html = """
  9. <!DOCTYPE html>
  10. <html>
  11.     <head>
  12.         <title>Chat</title>
  13.     </head>
  14.     <body>
  15.         <h1>WebSocket Chat</h1>
  16.         <h2>Your ID: <span id="ws-id"></span></h2>
  17.         <form action="" onsubmit="sendMessage(event)">
  18.             <input type="text" id="messageText" autocomplete="off"/>
  19.             <button>Send</button>
  20.         </form>
  21.         <ul id='messages'>
  22.         </ul>
  23.         <script>
  24.             var client_id = Date.now()
  25.             document.querySelector("#ws-id").textContent = client_id;
  26.             var ws = new WebSocket(`ws://localhost:8000/ws/${client_id}`);
  27.             ws.onmessage = function(event) {
  28.                 var messages = document.getElementById('messages')
  29.                 var message = document.createElement('li')
  30.                 var content = document.createTextNode(event.data)
  31.                 message.appendChild(content)
  32.                 messages.appendChild(message)
  33.             };
  34.             function sendMessage(event) {
  35.                 var input = document.getElementById("messageText")
  36.                 ws.send(input.value)
  37.                 input.value = ''
  38.                 event.preventDefault()
  39.             }
  40.         </script>
  41.     </body>
  42. </html>
  43. """
  46. class ConnectionManager:
  47.     def __init__(self):
  48.         self.active_connections: List[WebSocket] = []
  50.     async def connect(self, websocket: WebSocket):
  51.         await websocket.accept()
  52.         self.active_connections.append(websocket)
  54.     def disconnect(self, websocket: WebSocket):
  55.         self.active_connections.remove(websocket)
  57.     async def send_personal_message(self, message: str, websocket: WebSocket):
  58.         await websocket.send_text(message)
  60.     async def broadcast(self, message: str):
  61.         for connection in self.active_connections:
  62.             await connection.send_text(message)
  65. manager = ConnectionManager()
  68. @app.get("/")
  69. async def get():
  70.     return HTMLResponse(html)
  73. @app.websocket("/ws/{client_id}")
  74. async def websocket_endpoint(websocket: WebSocket, client_id: int):
  75.     await manager.connect(websocket)
  76.     try:
  77.         while True:
  78.             data = await websocket.receive_text()
  79.             await manager.send_personal_message(f"You wrote: {data}", websocket)
  80.             await manager.broadcast(f"Client #{client_id} says: {data}")
  81.     except WebSocketDisconnect:
  82.         manager.disconnect(websocket)
  83.         await manager.broadcast(f"Client #{client_id} left the chat")

尝试以下操作：

  • 使用多个浏览器选项卡打开应用程序。
  • 从这些选项卡中发送消息。
  • 然后关闭其中一个选项卡。

这将引发 WebSocketDisconnect 异常，并且所有其他客户端都会收到类似以下的消息：

  1. Client #1596980209979 left the chat

Tip

上面的应用程序是一个最小和简单的示例，用于演示如何处理和向多个 WebSocket 连接广播消息。

但请记住，由于所有内容都在内存中以单个列表的形式处理，因此它只能在进程运行时工作，并且只能使用单个进程。

如果您需要与 FastAPI 集成更简单但更强大的功能，支持 Redis、PostgreSQL 或其他功能，请查看 encode/broadcaster。

更多信息

要了解更多选项，请查看 Starlette 的文档：