-
-
Save francbartoli/2532f8bd8249a4cefa32f9c17c886a4b to your computer and use it in GitHub Desktop.
[Websocket demo for fastapi] example of broadcast using websockets for fastapi #fastapi #websockets
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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 = """ | |
<!DOCTYPE html> | |
<html> | |
<head> | |
<title>Chat</title> | |
</head> | |
<body> | |
<h1>WebSocket Chat</h1> | |
<form action="" onsubmit="sendMessage(event)"> | |
<input type="text" id="messageText" autocomplete="off"/> | |
<button>Send</button> | |
</form> | |
<ul id='messages'> | |
</ul> | |
<script> | |
var ws = new WebSocket("ws://localhost:8000/ws"); | |
ws.onmessage = function(event) { | |
var messages = document.getElementById('messages') | |
var message = document.createElement('li') | |
var content = document.createTextNode(event.data) | |
message.appendChild(content) | |
messages.appendChild(message) | |
}; | |
function sendMessage(event) { | |
var input = document.getElementById("messageText") | |
ws.send(input.value) | |
input.value = '' | |
event.preventDefault() | |
} | |
</script> | |
</body> | |
</html> | |
""" | |
@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. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment