# From https://github.com/tiangolo/fastapi/issues/258
from typing import List
from fastapi import FastAPI
from starlette.responses import HTMLResponse
from starlette.websockets import WebSocket, WebSocketDisconnect
app = FastAPI()
html = """
Chat
WebSocket Chat
"""
@app.get("/")
async def get():
return HTMLResponse(html)
class Notifier:
def __init__(self):
self.connections: List[WebSocket] = []
self.generator = self.get_notification_generator()
async def get_notification_generator(self):
while True:
message = yield
await self._notify(message)
async def push(self, msg: str):
await self.generator.asend(msg)
async def connect(self, websocket: WebSocket):
await websocket.accept()
self.connections.append(websocket)
def remove(self, websocket: WebSocket):
self.connections.remove(websocket)
async def _notify(self, message: str):
living_connections = []
while len(self.connections) > 0:
# Looping like this is necessary in case a disconnection is handled
# during await websocket.send_text(message)
websocket = self.connections.pop()
await websocket.send_text(message)
living_connections.append(websocket)
self.connections = living_connections
notifier = Notifier()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await notifier.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await websocket.send_text(f"Message text was: {data}")
except WebSocketDisconnect:
notifier.remove(websocket)
@app.get("/push/{message}")
async def push_to_connected_websockets(message: str):
await notifier.push(f"! Push notification: {message} !")
@app.on_event("startup")
async def startup():
# Prime the push notification generator
await notifier.generator.asend(None)
If you want to test it out, copy the above to a file main.py and start the server:
uvicorn main:app --reload
Then, open a few tabs at http://localhost:8000/ and send some chat messages (this should work the same as the base tutorial app). Then open http://localhost:8000/push/hello%20world and you should receive a push notification in each of your open tabs showing the message hello world.