PmSocket for Python

What PmSocket does

PmSocket sits on top of the repository's custom server router. A decorated route runs normally first. If that route returns None, PmSocket upgrades the request to an aiohttp.web.WebSocketResponse, sends a connection control message, and begins dispatching incoming frames.

Quick start

Create one adapter per endpoint, register its route, then attach an async message handler. The companion browser client is covered in the JavaScript guide.

import aiohttp
import aiohttp_ws

server = aiohttp_ws.Server(host="127.0.0.1", port=8080)
connection_registry = set()
chat_socket = aiohttp_ws.PmSocket(server)

@chat_socket.route("/chat/ws", method="GET")
async def open_chat(request):
    # Validate query parameters, authenticate, or create connection state here.
    # None tells PmSocket to perform the WebSocket upgrade.
    return None

@chat_socket.handler
async def receive_chat_message(ws, msg):
    connection_registry.add(ws)
    if msg.type is aiohttp.WSMsgType.TEXT and not msg.data.startswith("::"):
        await ws.send_str(f"echo:{msg.data}")

@chat_socket.on_disconnect
async def remove_chat_connection(ws):
    connection_registry.discard(ws)

server.run()

Connection lifecycle

1. The custom server router invokes the function decorated with route().
2. If it returns a non-None value, that value is passed back to the normal HTTP response machinery.
3. If it returns None, PmSocket creates WebSocketResponse(heartbeat=5, compress=False) and prepares it.
4. The server sends the text frame ::connect::. The bundled browser client marks itself alive when it receives this frame.
5. Each incoming frame goes through handle_msg(), which handles control text and then calls the registered application handler.
6. When iteration ends, the async disconnect hook is awaited, if one is registered.
7. The prepared WebSocketResponse is returned.

The server enables an aiohttp heartbeat every five seconds and disables WebSocket compression. These settings are fixed by the current implementation.

API reference

PmSocket(server)

Creates an adapter attached to the custom server instance. The instance stores one application message handler and one optional disconnect hook.

route(route, method="*", cache_time=None, server_cache_time=None, pass_cookies=False, **kwargs)

Registers the decorated function through the server's routing layer. Both synchronous and asynchronous route functions are accepted.

The underlying router supports normal response values including an aiohttp response, HTML text, bytes, JSON dictionaries, and optional status codes. In this wrapper, only a literal None triggers the WebSocket upgrade.

handler(function) / @socket

Sets the single application handler. The function receives (ws, msg). It may be synchronous or asynchronous. Registering another handler replaces the previous one.

@chat_socket
async def receive(ws, msg):
    ...

# Equivalent registration style:
@chat_socket.handler
async def receive(ws, msg):
    ...

Both decorator forms return the PmSocket instance, so the decorated function name no longer refers to the original function. Use another name if you need to call that function directly.

on_disconnect(function)

Registers one cleanup hook receiving ws. It is always awaited after the WebSocket loop, so define it with async def. Registering another hook replaces the previous one.

await handle_msg(ws, msg)

Processes one aiohttp message. This is called by the route loop and normally should not be called by application code. Text control messages are handled first; the same message is then passed to the application handler.

Control protocol

Handler contracts

The msg argument

msg is an aiohttp.WSMessage. Inspect msg.type before interpreting msg.data.

The ws argument

ws is an aiohttp.web.WebSocketResponse. Common operations include await ws.send_str(text), await ws.send_json(value), await ws.send_bytes(data), and await ws.close(). State is available through properties such as ws.closed and ws.close_code.

Practical patterns

Authenticate before upgrading

Perform authentication in the route or in an application-owned before-request hook. Arbitrary flags such as auth=True do not authenticate a connection. Return an HTTP response on failure and None after validation succeeds. The following fragment assumes the application provides tokens.is_valid().

@chat_socket.route("/chat/ws", method="GET")
async def open_chat(request):
    token = request.query.get("token")
    if not token or not await tokens.is_valid(token):
        return {"error": "unauthorized"}, 401
    return None

Parse JSON safely

import json

@chat_socket.handler
async def receive(ws, msg):
    if msg.type is not aiohttp.WSMsgType.TEXT:
        return
    if msg.data.startswith("::"):
        return
    try:
        payload = json.loads(msg.data)
    except json.JSONDecodeError:
        await ws.send_json({"type": "error", "message": "Invalid JSON"})
        return
    await dispatch(payload, ws)

Track connections outside PmSocket

PmSocket does not maintain a client collection. The route receives only request; the socket is created after it returns. Register ws when the first application message arrives, as in the quick start, then remove it in on_disconnect. Make cleanup idempotent because a connection may already have been removed elsewhere.

Adapter, socket, and message objects

There are three distinct objects in a PmSocket application: the adapter you construct, the live aiohttp socket created for each accepted request, and each native aiohttp message received on that socket. The adapter is shared between its connections; application state stored directly on it is also shared.

Connection-specific state can be held in an application dictionary keyed by ws. The route runs before ws exists, so record the socket when its first application message arrives. Keep arbitrary business data out of reserved control strings.

Complete PmSocket workflows

1. Reply to a correlated JSON request

This server responds to the browser request example in the JavaScript guide. JSON encoding and the request_id field belong to this application protocol, not PmSocket itself.

import json
import aiohttp
from aiohttp_ws import Server, PmSocket

server = Server(host="127.0.0.1", port=8080)
adapter = PmSocket(server)

@adapter.route("/state/ws", method="GET")
async def accept(request):
    return None

@adapter.handler
async def handle(ws, msg):
    if msg.type != aiohttp.WSMsgType.TEXT or msg.data.startswith("::"):
        return
    try:
        payload = json.loads(msg.data)
    except json.JSONDecodeError:
        await ws.send_json({"error": "invalid JSON"})
        return
    if not isinstance(payload, dict):
        return
    if payload.get("operation") == "get_state":
        await ws.send_json({"request_id": payload.get("request_id"),
                            "type": "state", "data": {"ready": True}})

server.run()

2. Track joined clients and broadcast

Use this alternative handler with an adapter already created above. Clients join by sending the text join. This registry tracks joined clients, not sockets that have connected without sending that message.

import asyncio
import aiohttp

joined = set()

@adapter.handler
async def chat(ws, msg):
    if msg.type != aiohttp.WSMsgType.TEXT or msg.data.startswith("::"):
        return
    if msg.data == "join":
        joined.add(ws)
        await ws.send_str("joined")
        return
    if ws not in joined:
        return
    peers = [peer for peer in joined if not peer.closed]
    results = await asyncio.gather(
        *(peer.send_str(msg.data) for peer in peers),
        return_exceptions=True,
    )
    for peer, result in zip(peers, results):
        if isinstance(result, Exception):
            joined.discard(peer)

@adapter.on_disconnect
async def cleanup(ws):
    joined.discard(ws)

3. Apply a payload rule and close the session

Another alternative handler accepts binary frames up to an application limit. WebSocketResponse already has its own transport limits; this rule additionally defines what this application accepts. A socket callback's return value alone does not send a response.

import aiohttp

@adapter.handler
async def binary_ingest(ws, msg):
    if msg.type == aiohttp.WSMsgType.TEXT:
        if not msg.data.startswith("::"):
            await ws.send_json({"error": "send a binary frame"})
        return
    if msg.type != aiohttp.WSMsgType.BINARY:
        return
    if len(msg.data) > 65536:
        await ws.send_json({"error": "fragment exceeds 64 KiB"})
        await ws.close()
        return
    await ws.send_json({"accepted_bytes": len(msg.data)})

Errors and implementation caveats

ValueError: Duplicate route.

The custom server already contains an equivalent route object. Give the endpoint a unique path and registration.

TypeError: 'NoneType' object is not callable

No message handler was registered before a frame arrived.

Disconnect cleanup after errors

The disconnect hook runs after normal WebSocket iteration, outside a finally block. An exception from the message handler can bypass it; handle application errors inside the handler if cleanup is required.

Disconnect hook fails when awaited

The hook was defined as a synchronous function. Use async def.

Messages arrive before application state exists

The server sends ::connect:: immediately after preparation. Initialize connection state before returning None, or lazily in the message handler.

Unexpected control strings

::ping:: and ::disconnect:: are still forwarded to the application handler. Filter the reserved protocol strings.

Compression or heartbeat needs differ

The adapter currently hard-codes compress=False and heartbeat=5; change the implementation or use a direct aiohttp WebSocket route.