backend / protocols / websockets / 02_python_implementation.md

Python WebSocket Implementations

6 interview angles 7 min read source

Python WebSocket Implementations

Three main paths: the websockets library (low-level, asyncio), FastAPI / Starlette (web-framework integration), or Django Channels (Django’s async layer). Plus older python-socketio for the Socket.IO protocol.

The choices

Library Best for
websockets low-level WS server / client; standalone or behind a framework
FastAPI / Starlette new APIs that mix REST + WS
Django Channels Django apps that need WS
python-socketio clients of an existing Socket.IO server, or you need Socket.IO features
aiohttp aiohttp-based apps; supports both server and client

For new Python services: FastAPI + websockets underneath is the most common modern stack.

websockets library — minimal server

pip install websockets
import asyncio
from websockets.asyncio.server import serve

async def handler(websocket):
    async for message in websocket:
        await websocket.send(f"echo: {message}")

async def main():
    async with serve(handler, "0.0.0.0", 8765):
        await asyncio.get_event_loop().create_future()  # run forever

asyncio.run(main())

async for message in websocket yields each incoming frame (text or bytes). Connection closes when the iterator ends.

For one-off send/receive:

async def handler(websocket):
    msg = await websocket.recv()              # blocks until message arrives
    await websocket.send("hello back")

websockets library — client

from websockets.asyncio.client import connect

async def client():
    async with connect("wss://example.com/ws") as ws:
        await ws.send("hello")
        response = await ws.recv()
        print(response)

asyncio.run(client())

The library handles handshake, framing, masking, ping/pong. You see strings and bytes.

FastAPI WebSockets

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"echo: {data}")
    except WebSocketDisconnect:
        pass

accept() completes the handshake. After that, receive_text() / receive_bytes() / receive_json() block for incoming messages; send_*() writes back.

WebSocketDisconnect is raised when the client closes (or network drops).

To reject a connection:

await websocket.close(code=4001, reason="Auth required")

FastAPI with broadcast — connection manager

A common pattern for chat / pub-sub:

class ConnectionManager:
    def __init__(self):
        self.active: list[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.active.append(ws)

    def disconnect(self, ws: WebSocket):
        self.active.remove(ws)

    async def broadcast(self, message: str):
        # send to all; tolerate failures (closed connections)
        for ws in list(self.active):
            try:
                await ws.send_text(message)
            except Exception:
                self.active.remove(ws)

manager = ConnectionManager()

@app.websocket("/room/{room_id}")
async def room(websocket: WebSocket, room_id: str):
    await manager.connect(websocket)
    try:
        while True:
            msg = await websocket.receive_text()
            await manager.broadcast(f"[{room_id}] {msg}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)

Limitation: in-memory list lives on one server. For multi-server, use Redis pub/sub. See 07_scaling.md.

Subprotocols and headers in FastAPI

@app.websocket("/ws")
async def ws_endpoint(websocket: WebSocket):
    # Inspect headers / subprotocols before accepting
    requested = websocket.headers.get("sec-websocket-protocol", "").split(", ")
    if "chat.v2" in requested:
        await websocket.accept(subprotocol="chat.v2")
    else:
        await websocket.close(code=1002)
        return
    # ...

accept() accepts the optional subprotocol. The client sees it in the handshake response.

For initial auth headers (cookies, JWT), inspect before accept():

token = websocket.cookies.get("session") or websocket.query_params.get("token")
if not validate(token):
    await websocket.close(code=4001)
    return

Django Channels

pip install channels channels-redis

asgi.py:

import os
from django.core.asgi import get_asgi_application
from channels.routing import ProtocolTypeRouter, URLRouter
from channels.auth import AuthMiddlewareStack
from myapp.routing import websocket_urlpatterns

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "myproject.settings")

application = ProtocolTypeRouter({
    "http": get_asgi_application(),
    "websocket": AuthMiddlewareStack(URLRouter(websocket_urlpatterns)),
})

Consumer:

# myapp/consumers.py
from channels.generic.websocket import AsyncWebsocketConsumer
import json

class ChatConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        self.room_name = self.scope["url_route"]["kwargs"]["room"]
        await self.channel_layer.group_add(self.room_name, self.channel_name)
        await self.accept()

    async def disconnect(self, close_code):
        await self.channel_layer.group_discard(self.room_name, self.channel_name)

    async def receive(self, text_data):
        data = json.loads(text_data)
        await self.channel_layer.group_send(
            self.room_name,
            {"type": "chat.message", "message": data["message"]},
        )

    async def chat_message(self, event):
        await self.send(text_data=json.dumps({"message": event["message"]}))

Routing:

# myapp/routing.py
from django.urls import re_path
from . import consumers

websocket_urlpatterns = [
    re_path(r"ws/chat/(?P<room>\w+)/$", consumers.ChatConsumer.as_asgi()),
]

channel_layer is the cross-process pub/sub layer — backed by Redis in production. group_add / group_send handles broadcast across many server processes. The headache of “in-memory list doesn’t work across servers” is solved by Channels’ channel layer.

Settings:

CHANNEL_LAYERS = {
    "default": {
        "BACKEND": "channels_redis.core.RedisChannelLayer",
        "CONFIG": {"hosts": [("redis", 6379)]},
    },
}

Channels handles auth via Django’s session middleware (AuthMiddlewareStack) — self.scope["user"] is the authenticated Django User. See 04_authentication.md.

Run with Daphne or Uvicorn:

daphne myproject.asgi:application
# or
uvicorn myproject.asgi:application --workers 4

Socket.IO — different protocol

Socket.IO is NOT WebSocket. It’s a protocol on top of WebSocket (with HTTP long-polling fallback) that adds rooms, namespaces, automatic reconnection, and acknowledgments.

# pip install python-socketio
import socketio

sio = socketio.AsyncServer(async_mode="asgi", cors_allowed_origins="*")
app = socketio.ASGIApp(sio)

@sio.event
async def connect(sid, environ):
    print(f"connected: {sid}")

@sio.event
async def message(sid, data):
    print(f"got: {data}")
    await sio.emit("response", {"echo": data}, to=sid)

Use Socket.IO when:

  • Frontend already uses Socket.IO (client and server protocols must match).
  • You want built-in rooms, namespaces, reconnection.
  • You need HTTP long-polling fallback for restrictive networks.

Use raw WebSockets when:

  • You want the simpler protocol.
  • Cross-language clients (Socket.IO clients exist in many languages but raw WS is more universal).
  • Bandwidth-sensitive (Socket.IO has framing overhead).

ASGI server choice

For async Python web apps including WebSockets:

Server Notes
Uvicorn reference ASGI server, used standalone or under Gunicorn
Hypercorn similar; supports HTTP/2 and HTTP/3
Daphne the original ASGI server, from Channels
Granian newer, Rust-based, fast

For production: typically gunicorn -k uvicorn.workers.UvicornWorker app:app --workers N. Gunicorn manages worker processes; Uvicorn is the per-worker ASGI server.

WebSocket workloads are I/O-bound; fewer workers (1-2 per CPU core) often outperform more.

Handling JSON

WebSocket frames are text or bytes. JSON convention:

import json

await websocket.send_json({"type": "message", "content": "hello"})
data = await websocket.receive_json()

FastAPI / Starlette has send_json / receive_json built in. With websockets library:

await ws.send(json.dumps({"hello": "world"}))
data = json.loads(await ws.recv())

Define a message envelope:

{
  "type": "chat_message",
  "id": "abc-123",
  "data": { ... },
  "timestamp": "2024-01-15T10:30:00Z"
}

type discriminates message kinds; handlers dispatch on it. id for acknowledgments. Treat the protocol like a schema — version it (type: "v2.chat_message" or a top-level version field) so you can evolve without breaking old clients.

Testing WebSocket endpoints

from fastapi.testclient import TestClient

def test_websocket():
    with TestClient(app) as client:
        with client.websocket_connect("/ws") as websocket:
            websocket.send_text("hello")
            response = websocket.receive_text()
            assert response == "echo: hello"

The TestClient context manager handles the handshake. websocket_connect() returns a sync WebSocket-like object even though the server is async.

For Channels:

from channels.testing import WebsocketCommunicator

async def test_consumer():
    communicator = WebsocketCommunicator(ChatConsumer.as_asgi(), "/ws/chat/room/")
    connected, _ = await communicator.connect()
    assert connected
    await communicator.send_json_to({"message": "hi"})
    response = await communicator.receive_json_from()
    await communicator.disconnect()

Performance considerations

  • Use async: WebSocket connections are long-lived; sync (one thread per connection) doesn’t scale. Always async/await.
  • Backpressure: a slow client makes await ws.send(...) block. The send queue can fill. See 05_connection_management.md.
  • Connection count limits: typically 10k+ per process on modern hardware. File descriptor limits and per-connection memory matter. Tune ulimit -n.
  • CPU per message: text-frame parsing is fast; JSON serialization is slower; permessage-deflate compression takes CPU.

For 100k+ concurrent connections, run multiple worker processes (or multiple servers); see 07_scaling.md.

Common pitfalls

  • Sync handlers in async server: blocking call (requests.get, time.sleep) stalls the event loop. Use httpx.AsyncClient, asyncio.sleep.
  • Sharing an in-memory connection list across workers: only works for single-worker setups. Use Redis pub/sub for multi-worker.
  • Forgetting to accept(): connection hangs in handshake limbo.
  • Not handling WebSocketDisconnect: exception in handler crashes the connection; in-flight state may leak.
  • Not closing on auth failure: client connects, gets no response, eventually times out. Close with 4001 immediately.
  • No timeout on receive_text: idle client holds the connection forever. Set timeouts; pair with heartbeats.

Common interview confusions

  • “FastAPI WebSockets are different from websockets library.” — FastAPI uses Starlette’s WebSocket, which is a thin async wrapper. The underlying ASGI server (Uvicorn) uses websockets or wsproto. Same protocol, different abstractions.
  • “Channels is just for WebSockets.” — Channels is Django’s async layer; WebSockets are the main use case but it also handles other protocols (chat, IoT, background tasks).
  • “Socket.IO is just WebSocket with extras.” — different protocol on top of WebSocket. Clients and servers MUST use matching Socket.IO versions.

Interview angle

  • “How would you build a WebSocket server in Python?” — modern: FastAPI’s @app.websocket("/path") decorator with async handler. await websocket.accept() + receive/send loop, WebSocketDisconnect for cleanup. Under the hood: websockets library or wsproto, ASGI server (Uvicorn).
  • “FastAPI vs Django Channels for WebSockets?” — FastAPI for new async-first APIs. Django Channels when integrating with an existing Django app (re-uses Django’s auth, ORM, models). Channels has the channel layer (Redis-backed pub/sub) built in.
  • “How do you broadcast to all connections from one server?” — keep a list of active WebSockets, iterate and send. For multi-server: Redis pub/sub or Channels’ channel layer — connections subscribe to a topic; broadcasts publish to Redis; every server forwards to its subscribers.
  • “How do you test a WebSocket endpoint?” — FastAPI’s TestClient.websocket_connect() for sync tests; Channels has WebsocketCommunicator. Both spin up the handler in-process and let you send/receive.
  • “What’s the difference between WebSocket and Socket.IO in Python?” — WebSocket is the protocol. Socket.IO is a different protocol on top of it (adds rooms, reconnection, fallback to HTTP long-polling, acknowledgments). Use raw WS unless your stack requires Socket.IO compatibility.
  • “How do you handle JSON over WebSocket?” — define a message envelope (type, data, optional id); send as JSON via send_json / receive_json. Version the protocol so old clients keep working.