WebSockets are an excellent way to send real-time updates to the users of your website. In this blog post, we'll show you how to set up websockets in a web application based on Starlette.

What are WebSockets?

At a basic level, WebSockets represent a bidirectional connection between the client and the server, which means that both client and server are able to send messages to each other. When establishing such a connection, the client first opens a "normal" HTTP connection to the server, and then requests the server to "upgrade" the connection to a WebSocket connection.

The server may or may not accept this request to upgrade, depending on what software is running on the backend. In the Python ecosystem, Django supports WebSockets using the channels package which, at the time of this writing, is an official Django project but isn't included with the default Django distribution.

Starlette has supported WebSockets since v0.6.0. In this blog post, we'll see how we can write a simple WebSockets server using Starlette.

Starlette HTTP Endpoints

Let's start the discussion with a normal web server.

In a Starlette application, everything on the server is accessible through an Endpoint mapped on to a URL. An Endpoint can be thought as of a class that represents a user request. HTTP resources in particular are available through an HTTPEndpoint.

Let's quickly look at a very minimal Starlette application:

from starlette.applications import Starlette
from starlette.endpoints import HTTPEndpoint
from starlette.routing import Route
from starlette.requests import Request
from starlette.responses import Response

class HelloEndpoint(HTTPEndpoint):
    def get(self, request: Request):
        return Response("<h1>Hello, World!</h1>")

instance = Starlette(
    routes=(
        Route("/hello", HelloEndpoint, name="hello"),
    ),
)

Put the code above in a file called main.py and run it using uvicorn main:instance --port 8000. You should be able to access http://localhost:8000/hello in your browser and see the spectacular "Hello, World!" message written across the screen.

In this example, the HelloEndpoint is a "normal" HTTP endpoint. Since WebSocket connections work a bit differently, Starlette provides a special base class for writing such endpoints. Let's explore that next.

Starlette WebSocket Endpoints

As we mentioned before, WebSocket connects are slightly different from normal connections. For this reason, Starlette provides a different base Endpoint class for you to extend from, appropriately called WebSocketEndpoint.

Let's use this class to add another endpoint on the server in our previous example, which would respond to WebSocket connections.

import logging

from starlette.applications import Starlette
from starlette.endpoints import HTTPEndpoint, WebSocketEndpoint
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Route, WebSocketRoute
from starlette.websockets import WebSocket

logger = logging.getLogger(__name__)

class HelloEndpoint(HTTPEndpoint):
    def get(self, request: Request):
        return Response("<h1>Hello, World!</h1>")

class WSEndpoint(WebSocketEndpoint):
    encoding = "bytes"

    async def on_connect(self, websocket: WebSocket):
        await websocket.accept()

        logger.info(f"Connected: {websocket}")

    async def on_receive(self, websocket: WebSocket, data: bytes):
        await websocket.send_bytes(b"OK!")

        logger.info("websockets.received")

    async def on_disconnect(self, websocket: WebSocket, close_code: int):
        logger.info(f"Disconnected: {websocket}")

instance = Starlette(
    routes=(
        Route("/hello", HelloEndpoint, name="hello"),
        WebSocketRoute("/ws", WSEndpoint, name="ws"),
    ),
)

As you can see, the WSEndpoint class looks a bit different from our previous HTTPEndpoint-based class. One of the things that's different is that there are no get or post handler methods. Instead, we have on_connect, on_receive, and on_disconnect. These three methods roughly correspond to the lifecycle stage of a client WebSocket connection.

on_connect

The on_connect method is called by the framework whenever a new client connects to the server. It accepts one websocket parameter which represents the connection object itself. Call websocket.accept() for the server to accept the client's connection request.

on_receive

The on_receive method is where the main communication happens. Whenever the client sends the server any data, this method is called along with the sent data passed as an argument.

In the example above, we ask the server to send back a simple "OK!" to the client as a generic response to anything that the client says.

on_disconnect

Starlette calls the on_disconnect method whenever the client/server connection closes. If you want the server to perform any cleanups after a client disconnects, this is the method to adjust.

A close_code argument is also passed to this method which contains the reason why the connection was closed. A complete list of all the possible close codes you might want to know about is available in the starlette.status module. This is helpful if you want to adjust the cleanup code depending on the nature of your application.

And that's basically it! In just a few more lines of code, we've built a production-ready WebSocket server that is able to communicate in real-time with multiple clients!

Testing WebSocket connections

Before closing out, let's quickly have a look at how to test the whole setup.

We could, of course, write some HTML and JavaScript code that connects to our backend over a WebSocket connection from the browser and send some messages back and forth between the client and the server to make sure everything is working as expected. And that would be a completely valid approach to testing.

Let's try something different. Let's try to learn Starlette a bit more and use its testing capabilities to make sure that our WebSocket server works as intended.

Starlette provides a TestClient class that for use as an HTTP client. As the name suggests, we normally use this class when writing application tests. We'll make an exception for this blog post and use this class to communicate with our WebSocket server.

from starlette.testclient import TestClient

from main import instance

def main():
    client = TestClient(instance)

    with client.websocket_connect("/ws") as websocket:
        websocket.send_bytes(b"Hello!")

        data = websocket.receive_bytes()

        print(data)

if __name__ == "__main__":
    main()

As we can see, this snippet instantiates a TestClient using the Starlette app instance we constructed in the main application file. We then use the websocket_connect context manager to get back a websocket object which we can use to communicate with the server. After the connection is established, the first thing we do is send the server a "Hello!" and then wait for the server to send something back. And in the next line, we print whatever we get back from the server.

Paste this code in a file called ws_test.py and run it using python ws_test.py. You should see output lines similar to the following:

~/W/p/g/code main ❯ python ws_test.py
[I 221012 22:03:46 ws:14] Connected: <starlette.websockets.WebSocket object at 0x1059a2850>
b'OK!'
[I 221012 22:03:46 ws:25] Disconnected: <starlette.websockets.WebSocket object at 0x1059a2850>

The first and last lines in the output are log messages from the calls to logger.info we made in the endpoint definition. And the second line is the "OK!" message that the server sends us back.

Conclusion

WebSockets are a very useful piece of technology that can help you enable rich end-user experiences.

The ability to push data from the server to the client instead of the client always having to ask the server for data is something very powerful. Building chat applications is a classic use-case that WebSockets enable, but the number of use-cases to which WebSockets apply is potentially limitless.

Starlette is one Python framework which supports this technology natively. So if you're building a SaaS, this is one more reason why you should consider using Starlette!

All code you see in this article is freely available on our Github: https://github.com/geniepy/snippets/tree/main/blog/starlette-websockets.


Photo by Christopher Robin Ebbinghaus on Unsplash

Launch your next SaaS quickly using Python 🐍

User authentication, Stripe payments, SEO-optimized blog, and many more features work out of the box on day one. Simply download the codebase and start building.