`WebSocketEndpoint`

Provide a class-based websocket handler.

from starlette.endpoints import WebSocketEndpoint


class App(WebSocketEndpoint):
    async def on_connect(self, websocket, **kwargs):
        ...
    async def on_receive(self, websocket, data):
        ...
    async def on_disconnect(self, websocket):
        ...

[jordaneremieff]

I’ve started on an implementation for this, it appears to work as expected in my test application. I’m happy to open a PR with tests, documentation, etc. after getting some feedback here.

How does this look so far?

class WebSocketEndpoint:
    def __init__(self, scope: Scope):
        self.scope = scope

    async def __call__(self, receive: Receive, send: Send):
        websocket = WebSocket(self.scope, receive=receive, send=send)
        kwargs = self.scope.get("kwargs", {})
        await self.on_connect(websocket, **kwargs)
        try:
            while True:
                message = await websocket.receive()
                websocket._raise_on_disconnect(message)
                if message["type"] == "websocket.receive":
                    await self.on_receive(websocket, message)
                elif message["type"] == "websocket.disconnect":
                    await self.on_disconnect(websocket)
        except WebSocketDisconnect as exc:
            await self.on_disconnect(websocket)

    async def on_connect(self, websocket, **kwargs):
        """Override to handle an incoming websocket connection"""
        await websocket.accept()

    async def on_receive(self, websocket, message):
        """Override to handle an incoming websocket message"""

    async def on_disconnect(self, websocket):
        """Override to handle a disconnecting websocket"""
        await websocket.close(code=1000)

[vlori2k]

I was thinking again about how we are passing the WebSocket instance throughout the WebSocketEndpoint instance in the current proposed pattern, and I was wondering if it may be better to store the websocket state on the endpoint instead as @tomchristie initially suggested. This also made me wonder if there is any reason not to inherit from the WebSocket class and expose its methods directly.

Here are two potential ways of implementing that seem to be a bit more flexible based on how I’ve been using it:

class WebSocketEndpoint(WebSocket):

    encoding = None

    async def __call__(self, receive: Receive, send: Send):
        self._receive = receive
        self._send = send
        close_code = 1000

        try:
            while True:
                message = await self.receive()
                if message["type"] == "websocket.connect":
                    await self.on_connect(message)
                elif message["type"] == "websocket.receive":
                    data = await self.decode(message)
                    await self.on_receive(data)
                elif message["type"] == "websocket.disconnect":
                    close_code = message["code"]
                    return
        finally:
            if close_code == 1006:
                raise RuntimeError("WebSocket error, connection closed abnormally")
            await self.on_disconnect(close_code=close_code)

    async def decode(self, message):
        if self.encoding == "text":
            if "text" not in message:
                await self.close(code=1003)
                raise RuntimeError("Expected text websocket messages, but got bytes")
            return message["text"]

        elif self.encoding == "bytes":
            if "bytes" not in message:
                await self.close(code=1003)
                raise RuntimeError("Expected bytes websocket messages, but got text")
            return message["bytes"]

        elif self.encoding == "json":
            if "bytes" not in message:
                await self.close(code=1003)
                raise RuntimeError(
                    "Expected JSON to be transferred as bytes websocket messages, but got text"
                )
            return json.loads(message["bytes"].decode("utf-8"))

        assert (
            self.encoding is None
        ), f"Unsupported 'encoding' attribute {self.encoding}"

        return message["text"] if "text" in message else message["bytes"]

    async def on_connect(self, message):
        """Override to handle an incoming websocket connection"""
        await self.accept()

    async def on_receive(self, data):
        """Override to handle an incoming websocket message"""
        await self.receive(data)

    async def on_disconnect(self, close_code):
        """Override to handle a disconnecting websocket"""
        await self.close(code=close_code)

Alternatively, we could do the same as above except like this:

class WebSocketEndpoint:
    def __init__(self, scope: Scope) -> None:
        self.scope = scope

    async def __call__(self, receive: Receive, send: Send):
        self.websocket = WebSocket(self.scope, receive=receive, send=send)
        ...

    async def on_connect(self, message):
        """Override to handle an incoming websocket connection"""
        await self.websocket.accept()

    async def on_receive(self, data):
        """Override to handle an incoming websocket message"""
        await self.websocket.receive(data)

    async def on_disconnect(self, close_code):
        """Override to handle a disconnecting websocket"""
        await self.websocket.close(code=close_code)

Likely a bit of finessing required. Also, I haven’t been across all of the discussion surrounding the disconnect/close behaviour, so maybe there is something from there that could help inform this as well.

Thoughts?

Does this work? im having trouble with " RuntimeError: Unexpected ASGI message ‘websocket.send’, after sending ‘websocket.close’."