Backend / Protocols / WebSockets / 02_python_implementation.md

Python WebSocket Implementations

Updated 6 interview angles 7 min read source
On this page15
  1. The choices
  2. websockets library — minimal server
  3. websockets library — client
  4. FastAPI WebSockets
  5. FastAPI with broadcast — connection manager
  6. Subprotocols and headers in FastAPI
  7. Django Channels
  8. Socket.IO — different protocol
  9. ASGI server choice
  10. Handling JSON
  11. Testing WebSocket endpoints
  12. Performance considerations
  13. Common pitfalls
  14. Common interview confusions
  15. Interview angle

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

bash
pip install websockets
python
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):
        # run forever
        await asyncio.get_event_loop().create_future()

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:

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

websockets library — client

python
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

python
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:

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

FastAPI with broadcast — connection manager

A common pattern for chat / pub-sub:

python
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 Scaling WebSocket Applications.

Subprotocols and headers in FastAPI

python
@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():

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

Django Channels

bash
pip install channels channels-redis

asgi.py:

python
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:

python
# 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:

python
# 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:

python
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 WebSocket Authentication.

Run with Daphne or Uvicorn:

bash
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.

python
# 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:

python
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:

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

Define a message envelope:

json
{
  "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

python
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:

python
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 WebSocket Connection Management.
  • 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 Scaling WebSocket Applications.

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 6

  • “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.