Skip to content

Instantly share code, notes, and snippets.

@francbartoli
Forked from wshayes/demo.py
Created June 2, 2019 12:48
Show Gist options
  • Save francbartoli/2532f8bd8249a4cefa32f9c17c886a4b to your computer and use it in GitHub Desktop.
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
# 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