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.