Created
June 8, 2019 22:17
-
-
Save wshayes/8b0bf51397131c018d0e3f735eb02784 to your computer and use it in GitHub Desktop.
[FastAPI Websocket Example #2] #fastapi
This file contains hidden or 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
# 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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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?