-
-
Save wshayes/8b0bf51397131c018d0e3f735eb02784 to your computer and use it in GitHub Desktop.
# https://github.com/tiangolo/fastapi/issues/258#issuecomment-495975801 | |
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) |
Hi - I grabbed this from the FastAPI issue queue so I didn't write it, but this basically starts a task on the event loop in FastAPI that is always checking for WebSocket messages from the connected clients.
async def get_notification_generator(self): while True: message = yield await self._notify(message)
Can you explain what's happening here? Never seen something like this before, though am not an expert at async programming.
Hi,
Can you explain what is "generator.asend(None)"? I don't see the function declared anywhere
Is the "asend" a typo ?
Sorry this is so late. asend
is an asynchronous send - https://www.python.org/dev/peps/pep-0525/#why-the-asend-and-athrow-methods-are-necessary
Okay so from what I can understand, this is a bit like goroutines in Go. message = yield
will block till generator.asend is called and subsequently the arg passed in asend will be received inside the message variable, am I correct?
async def get_notification_generator(self): while True: message = yield await self._notify(message)
Can you explain what's happening here? Never seen something like this before, though am not an expert at async programming.