WebSockets in Django 3.1

Alex Oleshkevich — Sat, 08 Aug 2020 22:10:47 +0000

This is a repost of my post on medium.com

Together with my wife, we run a small digital agency. We use Django as a primary web development framework and love simplicity.

In this post, I will guide you on how to enable WebSockets in your Django application without installing third-party apps.

Django has introduced the ASGI interface since version 3.0 and async views in 3.1. Our solution will be based on async views. In this tutorial, we will use Python 3.7 and Django 3.1.

Introduction into WebSockets ASGI interface

ASGI is a replacement protocol to good-old WSGI protocol that served us for years and it will become a de-facto standard in Python web frameworks in the next 2–3 years.

So, how does WebSocket work in this context? Let us find it!

The communication between WebSocket clients and your application is event-based. The ASGI specification defines two types of events: send and receive.

Receive events. These are events that clients send to your application. Let's look at them:

websocket.connect is sent when the client tries to establish a connection with our application
websocket.receive is sent when the client sends data to our app
websocket.disconnect tells us that the client has disconnected.

Send events are emitted by our application to a client (e.g. a browser). Here is a list of them:

websocket.accept — we send this event back to the client if we want to allow the connection
websocket.send — with this event, we push data to the client
websocket.close is emitted by the application when we want to abort the connection.

Now, as we know all participants of that party, it is the time to speak about their order.

When a browser opens a connection, the ASGI protocol server (we will talk about this later) sends us websocket.connect event. Our application must respond to it with either websocket.accept or websocket.close according to our logic. It is simple: emit websocket.accept if you allow the connection or emit websocket.close to cancel the connection. You may want to cancel the connection, for example, if the user has no permissions to connect or is not logged in. I will assume that you allow the connection in the next steps.

After you accepted the connection, the application is ready to send and receive data via that socket using websocket.send and websocket.receive events.

Finally, when the browser leaves your page or refreshes it, a websocket.disconnect is sent to the application. As a developer, you still have control over the connection and can abort the connection by sending websocket.close event at any time you wish.

This was a brief description of how ASGI processes WebSockets. It is not scoped to Django, it works for any other ASGI compatible web framework like Starlette or FastAPI.

Setting up Django apps

In this tutorial, I am not going to cover Django installation and setup topics. Also, I assume that you have Django installed and operating.

First, we have to create a new Django application. This app will keep the custom URL pattern function, ASGI middleware, and WebSocket connection class.

Let's make a new app using this command:

django-admin startapp websocket

Okay, now let's make a new little helper function for developer convenience. This function will be a simple alias for path function at the moment.

Add urls.py to the websocket app with this content:

from django.urls import path

websocket = path

Now you can configure WebSocket URLs in a distinct way. Time to create your first WebSocket view! To keep things simple and reusable we will make another Django app named, say, users. Don’t forget to enable both applications in the INSTALLED_APPS setting!

django-admin startapp users

Implementing ASGI middleware

The middleware will be our glue code between WebSockets and asynchronous views provided by Django. The middleware will intercept WebSocket requests and will dispatch them separately from the Django default request handler. When you created a new Django project, the installed has added a new file named asgi.py to the project installation directory. You will find the ASGI application in it. This is the application we are going to use instead of one defined in wsgi.py.

Create a new websocket/middleware.py file and put the code in it:

from django.urls import resolve
from .connection import WebSocket

def websockets(app):
    async def asgi(scope, receive, send):
        if scope["type"] == "websocket":
            match = resolve(scope["raw_path"])
            await match.func(WebSocket(scope, receive, send), *match.args, **match.kwargs)
            return
    await app(scope, receive, send)
return asgi

Every ASGI middleware is a callable that accepts another callable. In the middleware, we test if the request type is websocket, and if so, we call Django’s URL resolver for a dispatchable view function. By the way, a 404 error will be raised if the resolver fails to find a view matching the URL.

Now, open project_name/asgi.py file and wrap default application with this middleware:

from django.core.asgi import get_asgi_application
from websocket.middleware import websockets
application = get_asgi_application()
application = websockets(application)

Since that moment, every request made will be caught by our middleware and tested for its type. If the type is websocket then the middleware will try to resolve and call a view function.

Don’t mind at the moment about missing import from the .connection module. We are about to make it in a minute.

Add WebSocket connection

The role of the WebSocket connection is similar to the request object you use in your views. The connection will encapsulate request information along with methods that assist you in receiving and sending the data. This connection will be passed as the first argument of our WebSocket view functions.

Create websocket/connection.py with the contents from the gist below. To make life easier we will also enumerate all possible WebSocket events in classes, add Headers class to access request headers, and QueryParams to get variables from a query string.

import json
import typing as t
from urllib import parse


class State:
    CONNECTING = 1
    CONNECTED = 2
    DISCONNECTED = 3


class SendEvent:
    """Lists events that application can send.
    ACCEPT - Sent by the application when it wishes to accept an incoming connection.
    SEND - Sent by the application to send a data message to the client.
    CLOSE - Sent by the application to tell the server to close the connection.
        If this is sent before the socket is accepted, the server must close
        the connection with a HTTP 403 error code (Forbidden), and not complete
        the WebSocket handshake; this may present on some browsers as 
        a different WebSocket error code (such as 1006, Abnormal Closure).
    """

    ACCEPT = "websocket.accept"
    SEND = "websocket.send"
    CLOSE = "websocket.close"


class ReceiveEvent:
    """Enumerates events that application can receive from protocol server.
    CONNECT - Sent to the application when the client initially 
        opens  a connection and is about to finish the WebSocket handshake.
        This message must be responded to with either an Accept message or a Close message 
        before the socket will pass websocket.receive messages.
    RECEIVE - Sent to the application when a data message is received from the client.
    DISCONNECT - Sent to the application when either connection to the client is lost, 
        either from the client closing the connection, 
        the server closing the connection, or loss of the socket.
    """

    CONNECT = "websocket.connect"
    RECEIVE = "websocket.receive"
    DISCONNECT = "websocket.disconnect"


class Headers:
    def __init__(self, scope):
        self._scope = scope

    def keys(self):
        return [header[0].decode() for header in self._scope["headers"]]

    def as_dict(self) -> dict:
        return {h[0].decode(): h[1].decode() for h in self._scope["headers"]}

    def __getitem__(self, item: str) -> str:
        return self.as_dict()[item.lower()]

    def __repr__(self) -> str:
        return str(dict(self))


class QueryParams:
    def __init__(self, query_string: str):
        self._dict = dict(parse.parse_qsl(query_string))

    def keys(self):
        return self._dict.keys()

    def get(self, item, default=None):
        return self._dict.get(item, default)

    def __getitem__(self, item: str):
        return self._dict[item]

    def __repr__(self) -> str:
        return str(dict(self))


class WebSocket:
    def __init__(self, scope, receive, send):
        self._scope = scope
        self._receive = receive
        self._send = send
        self._client_state = State.CONNECTING
        self._app_state = State.CONNECTING

    @property
    def headers(self):
        return Headers(self._scope)

    @property
    def scheme(self):
        return self._scope["scheme"]

    @property
    def path(self):
        return self._scope["path"]

    @property
    def query_params(self):
        return QueryParams(self._scope["query_string"].decode())

    @property
    def query_string(self) -> str:
        return self._scope["query_string"]

    @property
    def scope(self):
        return self._scope

    async def accept(self, subprotocol: str = None):
        """Accept connection.
        :param subprotocol: The subprotocol the server wishes to accept.
        :type subprotocol: str, optional
        """
        if self._client_state == State.CONNECTING:
            await self.receive()
        await self.send({"type": SendEvent.ACCEPT, "subprotocol": subprotocol})

    async def close(self, code: int = 1000):
        await self.send({"type": SendEvent.CLOSE, "code": code})

    async def send(self, message: t.Mapping):
        if self._app_state == State.DISCONNECTED:
            raise RuntimeError("WebSocket is disconnected.")

        if self._app_state == State.CONNECTING:
            assert message["type"] in {SendEvent.ACCEPT, SendEvent.CLOSE}, (
                'Could not write event "%s" into socket in connecting state.'
                % message["type"]
            )
            if message["type"] == SendEvent.CLOSE:
                self._app_state = State.DISCONNECTED
            else:
                self._app_state = State.CONNECTED

        elif self._app_state == State.CONNECTED:
            assert message["type"] in {SendEvent.SEND, SendEvent.CLOSE}, (
                'Connected socket can send "%s" and "%s" events, not "%s"'
                % (SendEvent.SEND, SendEvent.CLOSE, message["type"])
            )
            if message["type"] == SendEvent.CLOSE:
                self._app_state = State.DISCONNECTED

        await self._send(message)

    async def receive(self):
        if self._client_state == State.DISCONNECTED:
            raise RuntimeError("WebSocket is disconnected.")

        message = await self._receive()

        if self._client_state == State.CONNECTING:
            assert message["type"] == ReceiveEvent.CONNECT, (
                'WebSocket is in connecting state but received "%s" event'
                % message["type"]
            )
            self._client_state = State.CONNECTED

        elif self._client_state == State.CONNECTED:
            assert message["type"] in {ReceiveEvent.RECEIVE, ReceiveEvent.DISCONNECT}, (
                'WebSocket is connected but received invalid event "%s".'
                % message["type"]
            )
            if message["type"] == ReceiveEvent.DISCONNECT:
                self._client_state = State.DISCONNECTED

        return message

    async def receive_json(self) -> t.Any:
        message = await self.receive()
        self._test_if_can_receive(message)
        return json.loads(message["text"])

    async def receive_jsonb(self) -> t.Any:
        message = await self.receive()
        self._test_if_can_receive(message)
        return json.loads(message["bytes"].decode())

    async def receive_text(self) -> str:
        message = await self.receive()
        self._test_if_can_receive(message)
        return message["text"]

    async def receive_bytes(self) -> bytes:
        message = await self.receive()
        self._test_if_can_receive(message)
        return message["bytes"]

    async def send_json(self, data: t.Any, **dump_kwargs):
        data = json.dumps(data, **dump_kwargs)
        await self.send({"type": SendEvent.SEND, "text": data})

    async def send_jsonb(self, data: t.Any, **dump_kwargs):
        data = json.dumps(data, **dump_kwargs)
        await self.send({"type": SendEvent.SEND, "bytes": data.encode()})

    async def send_text(self, text: str):
        await self.send({"type": SendEvent.SEND, "text": text})

    async def send_bytes(self, text: t.Union[str, bytes]):
        if isinstance(text, str):
            text = text.encode()
        await self.send({"type": SendEvent.SEND, "bytes": text})

    def _test_if_can_receive(self, message: t.Mapping):
        assert message["type"] == ReceiveEvent.RECEIVE, (
            'Invalid message type "%s". Was connection accepted?' % message["type"]
        )

Add your first WebSocket view

Our project is set up to handle WebSocket connections. The only thing left is a WebSocket view function. We would also need a template view to serve an HTML page.

# users/views.py
from django.views.generic.base import TemplateView
class IndexView(TemplateView):
    template_name = "index.html"
async def websocket_view(socket):
    await socket.accept()
    await socket.send_text('hello')
    await socket.close()

Mount both views in the root urls.py

# project_name/urls.py
from django.urls import path
from websocket.urls import websocket
from users import views
urlpatterns = [
    path("", views.IndexView.as_view()),
    websocket("ws/", views.websocket_view),
]

users/templates/index.html should contain this script:

<script>
    new WebSocket('ws://localhost:8000/ws/');
</script>

This is a bare minimum to establish a WebSocket connection.

Start the development server

The Django’s runserver command does not use application defined in asgi.py at the time of writing this post. We need to use a 3rd-party application server. I will use Uvicorn.

pip install uvicorn

Once installed start the server passing ASGI application as the first positional argument:

uvicorn project_name.asgi:application --reload --debug

Navigate to http://localhost:8000/, open browser’s console, switch to Network tab and observe the WebSockets working.

Echo server

The WebSocket view we created is useless. It sends one message and then closes the connection. We will refactor it in a simple echo server that replies to a client using the incoming message text.

Replace websocket_view in users/views.py with this code:

async def websocket_view(socket: WebSocket):
    await socket.accept()
    while True:
        message = await socket.receive_text()
        await socket.send_text(message)

and replace contents of users/templates/index.html with this:

<script>
let socket = new WebSocket('ws://localhost:8000/ws/');
let timer = null;
socket.onopen = () => {
    timer = setInterval(() => {
        socket.send('hello');
    }, 1000);
};
socket.onclose = socket.onerror = () => {
    clearInterval(timer);
};
</script>

The updated code will send hello text to our application every one second and our application will respond to it with the same message.

Conclusion

In this post, I demonstrated how to add WebSocket support to your Django project using only Django 3.1 and the standard python library. Yes, I know that Uvicorn still needs to be installed but this is a limitation of the Django dev server at the moment.

DEV Community: Alex Oleshkevich