Setup A Python Flask Web Server With WebSockets

Overview

When you search flask websockets the first thing that comes up is Flask-SocketIO. Turns on, SocketIO and WebSockets aren’t the same thing. Took me a while to figure that out. Once I did, I put together this example file of how to get flask and websockets running together via a single script. It:

  • spins up a flask web server on port 5050
  • spins up websockets on port 5051
  • sets up the websockets so that incomings messages go to each client except the one that sent the original message
  • provides an index file at the root of the flask server that contains an example page (at http://localhost:5000/) that you can load multiple times to see the communication take place
  • Note that the way I’m excluding the client that sends the original message might not work all the time. I’ve seen it not send to other clients as well. Not sure about the behavior. If you have problems, see the note below about how to send to every client (including the one that sent the original message). (Maybe a single browser can use the same web socket for multiple tabs? If you know, let me know on Twitter)
#!/usr/bin/env python3

import asyncio
import os
import subprocess
import sys
import time
import websockets

from flask import Flask, request, Response

python_path = sys.executable
script_path = os.path.realpath(__file__) 
domain = 'localhost'
http_port = '5050'
ws_port = '5051'

### Flask Stuff
app = Flask(__name__)

html_page = '''<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>WebSocket Example</title>
</head>
<body>
<input id="textField" />
<p id="outputField" />
<script>
'''

html_page += f'websocket = new WebSocket("ws://{domain}:{ws_port}/");'

html_page += '''
textField = document.getElementById('textField');
outputField = document.getElementById('outputField');
textField.addEventListener('input', () => {
    currentText = textField.value;  
    websocket.send(currentText);
});

websocket.addEventListener('message', (event) => {
    outputField.innerHTML = event.data
})


</script>
</body>
</html> 
'''

@app.route('/', methods=['GET'])
def respond2():
    return Response(
        status=200,
        response=html_page
    )

### Websocket stuff 

CONNECTIONS =  set()

async def register(websocket):
    CONNECTIONS.add(websocket)

async def unregister(websocket):
    CONNECTIONS.remove(websocket)

async def notify_users(message, websocket):
    connection_list = []
    for connection in CONNECTIONS:
        if connection != websocket:
            connection_list.append(connection)

    await asyncio.wait([
        connection.send(message) for connection in connection_list
    ])

async def message_control(websocket, path):
    await register(websocket)
    try:
        await websocket.send("Connected")
        async for message in websocket:
            print(message)
            await notify_users(message, websocket)
    finally:
        await unregister(websocket)


if len(sys.argv) == 2: 
    if sys.argv[1] == 'run_socket':
        print("\nStarting websocker server")
        start_server = websockets.serve(message_control, domain, ws_port)

        asyncio.get_event_loop().run_until_complete(start_server)
        asyncio.get_event_loop().run_forever()

else:
    time.sleep(1)
    print("Triggering websocket server")
    return_value = (subprocess.Popen([
        python_path,
        script_path,
        'run_socket'
    ]))

    print("Starting flask server")
    app.run(port=http_port)

If you want the messages to be sent to all clients (including the one that originally sent the message), then change this:

async def notify_users(message, websocket):
    connection_list = []
    for connection in CONNECTIONS:
        if connection != websocket:
            connection_list.append(connection)

    await asyncio.wait([
        connection.send(message) for connection in connection_list
    ])

to this:

async def notify_users(message, websocket):
    connection_list = []
    for connection in CONNECTIONS:
        connection_list.append(connection)

    await asyncio.wait([
        connection.send(message) for connection in connection_list
    ])

That’s it. Happy Web Socketing!

2021-05-29T16:35:21-0400

Random Related Links