Add a WebSocket Route to your Flask 2.x Application

The WebSocket protocol was standardized 10 years ago (in 2011, can you believe it?) so I'm sure for many of it you does not need an introduction. But in case you are not familiar with it, WebSocket is an extension to the HTTP protocol that provides a permanent, bi-directional communication channel between a client and the server, where both sides can send and receive data in real time, free of the constraints of the request/response cycle of HTTP.

Flask, being a minimalist web framework, does not have WebSocket support built-in. The old Flask-Sockets extension, which has not been maintained in the last few years, provided support for this protocol. My own Flask-SocketIO extension has also been a favorite, but this is Socket.IO, which is a more complex protocol that combines WebSocket and HTTP.

The good news is that if you are using Flask 2 you now have a new extension (also created by myself) called Flask-Sock, which provides modern WebSocket support for your application. In this article I'm going to show you how to work with this extension.

This article was voted by my supporters on Patreon. Would you like to support my work, and as a thank you be able to vote on my future articles and also have access to a chat room where I hang out? Become a Patron!

Introducing the Flask-Sock Extension

The Flask-Sock extension is installed with pip:

pip install flask-sock

The extension is added to your Flask application by creating a Sock() instance. If you have a global app object, you can use the direct initialization method:

sock = Sock(app)

If instead you create your application instance in a factory function, then the two-step initialization method works as well:

sock = Sock()

def create_app():
    app = Flask(__name__)
    sock.init_app(app)

The sock instance has a route decorator, that works pretty much like Flask's own, but it adds the WebSocket protocol handshake so that the route can speak WebSocket instead of HTTP. Here is a simple server that echoes back to the client anything it sends:

@sock.route('/echo')
def echo(ws):
    while True:
        data = ws.receive()
        ws.send(data)

The ws object that is passed to the route is the actual WebSocket connection, which the function can use to exchange information with the client through the send() and receive() methods. It goes without saying that when multiple clients are connected at the same time, each gets its own ws object, and the communication between the server and each client is private.

Looking at the code example above, the main difference between a Flask route and a WebSocket route is that the WebSocket function is designed to run for a long time. WebSocket connections often last for as long as the user has a page in the browser open, so it may be several minutes or even hours. Unlike standard HTTP routes, your WebSocket handler will very likely have to be implemented with a loop, so that it remains active for as long as it is needed. The WebSocket function is not expected to return a response, the Flask-Sock extension takes care of generating the response as required by the WebSocket protocol.

How do you end a WebSocket connection? The most common situation is that the user navigates away from the page that created the connection, At this point the browser will close the WebSocket, and from then on the ws.receive() and ws.send() methods are going to raise a ConnectionClosed exception, which is automatically handled by the Flask-Sock extension to end the route gracefully. You can catch this exception in the route to perform custom cleanup if necessary.

A second way to end the WebSocket connection occurs when the server initiates the close. To do this, the WebSocket route needs to call ws.close(), or alternatively, just exit the function.

The ws.route decorator is fully integrated with Flask, and in fact, it uses the app.route decorator internally. In particular, all these things work:

• Variable components in the URL. If your URL has variables, the ws argument must be included before those variables in the function definition.
• Additional decorators. You can add any decorators that are designed to work with Flask routes, such as Flask-Login's login_required and any others.
• Before and after request handlers. These are invoked for WebSocket routes as well.
• current_app, request, g and session. These work the same as in Flask routes, with one exception regarding the session variable, explained below.

An important feature of HTTP that is not available in WebSocket is the ability to set cookies. In HTTP, cookies are set in the response headers. WebSocket does not use an HTTP response, so it is not possible to set a cookie when the WebSocket connection ends. This means that changes made to the session object in the WebSocket handler function are not saved.

Running your Flask + WebSocket server

To run your WebSocket enabled web server during development, just run your application in the way you normally do. Both the flask run and app.run() methods of running the server are compatible with Flask-Sock.

To run your Flask + WebSocket server in production you can use Gunicorn. You will normally want to enable multithreading, because as discussed earlier, a WebSocket route will run for a long time, and without threads each WebSocket connection would consume an entire worker. The following example creates a Gunicorn server with four workers, each using 100 threads, allowing a maximum of 400 active WebSocket connections (though you will always want to leave some threads available to handle other HTTP requests):

gunicorn -b 0.0.0.0:5000 --workers 4 --threads 100 module:app

If you prefer to use a greenlet based asynchronous web server, Flask-Sock is also compatible with eventlet and gevent. The documentation has all the details. The uWSGI web server is currently not supported.

A Complete Example

Let's have a look at a complete example. Here is the Flask application code:

from flask import Flask, render_template
from flask_sock import Sock

app = Flask(__name__)
sock = Sock(app)


@app.route('/')
def index():
    return render_template('index.html')


@sock.route('/echo')
def echo(sock):
    while True:
        data = sock.receive()
        sock.send(data)

This is based in the echo server code shown above. The application has a root URL that returns the index.html client page, and a /echo route that implements a WebSocket echo endpoint.

Here is the index.html page, which you have to save in a templates subdirectory:

<!doctype html>
<html>
  <head>
    <title>Flask-Sock Demo</title>
  </head>
  <body>
    <h1>Flask-Sock Demo</h1>
    <div id="log"></div>
    <br>
    <form id="form">
      <label for="text">Input: </label>
      <input type="text" id="text" autofocus>
    </form>
    <script>
      const log = (text, color) => {
        document.getElementById('log').innerHTML += `<span style="color: ${color}">${text}</span><br>`;
      };

      const socket = new WebSocket('ws://' + location.host + '/echo');
      socket.addEventListener('message', ev => {
        log('<<< ' + ev.data, 'blue');
      });
      document.getElementById('form').onsubmit = ev => {
        ev.preventDefault();
        const textField = document.getElementById('text');
        log('>>> ' + textField.value, 'red');
        socket.send(textField.value);
        textField.value = '';
      };
    </script>
  </body>
</html>

The HTML page includes a <script> section with the client-side logic. The log() auxiliary function writes a string to the page, using the given color. This is used to write outgoing data in red, and incoming data in blue, as you can see in the screenshot below:

The application uses the WebSocket class provided by the browser to make a connection to the /echo endpoint. It then installs a listener for the message event, which fires whenever the server sends data through the connection. The handler prints the data to the page. The page has a form with an <input> element that accepts text from the user. A second event handler is defined on the submit event for the form. The handler for this event logs the text entered by the user to the page, and then writes it to the WebSocket connection, using its send() method.

To try out this application, save the Python file as app.py and the HTML file as index.html in a templates subdirectory. Then run the application with flask run. Connect with your browser at the http://localhost:5000 URL to use the application.

What Flask-Sock Doesn't Do

Flask-Sock is designed to be a simple extension to implement WebSocket communication between a client and the server. It is not a goal of this extension to implement advanced workflows involving the broadcasting of messages to all or to subsets of clients. Your application can implement this if desired, but if you need something of this complexity you may be better off using Flask-SocketIO. I have written an article about it too!