WebSockets#
Handling WebSocket in an application often involves dealing with low level constructs
such as the socket itself, setting up a loop and listening for incoming data, handling
exceptions, and parsing incoming and serializing outgoing data. In addition to the
low-level websocket route handler, Litestar offers two
high level interfaces:
These treat a WebSocket handler like any other route handler; As a callable that takes in incoming data in an already pre-processed form and returns data to be serialized and sent over the connection. The low level details will be handled behind the curtains.
from litestar import Litestar
from litestar.handlers.websocket_handlers import websocket_listener
@websocket_listener("/")
async def handler(data: str) -> str:
return data
app = Litestar([handler])
This handler will accept connections on /, and wait to receive data. Once a message
has been received, it will be passed into the handler function defined, via the data
parameter. This works like a regular route handler, so it’s possible to specify the
type of data which should be received, and it will be converted accordingly.
Note
Contrary to websocket route handlers, function decorated with
websocket_listener don’t have to be
asynchronous.
Receiving data#
Data can be received in listener via the data parameter. The data passed to this
will be converted / parsed according to the given type annotation and supports
str, bytes, or arbitrary dicts / or lists in the
form of JSON.
from typing import Dict
from litestar import Litestar, websocket_listener
@websocket_listener("/")
async def handler(data: Dict[str, str]) -> Dict[str, str]:
return data
app = Litestar([handler])
from litestar import Litestar, websocket_listener
@websocket_listener("/")
async def handler(data: dict[str, str]) -> dict[str, str]:
return data
app = Litestar([handler])
from litestar import Litestar, websocket_listener
@websocket_listener("/")
async def handler(data: str) -> str:
return data
app = Litestar([handler])
from litestar import Litestar, websocket_listener
@websocket_listener("/")
async def handler(data: bytes) -> str:
return data.decode("utf-8")
app = Litestar([handler])
Important
Contrary to route handlers, JSON data will only be parsed but not validated. This is a currently limitation of the implementation and will change in future versions.
Sending data#
Sending data is done by simply returning the value to be sent from the handler function.
Similar to receiving data, type annotations configure how the data is being handled.
Values are are not str or bytes are assumed to be JSON encodable and
will be serialized accordingly before being sent. This serialization is available for
all data types currently supported by Litestar (
dataclasses, TypedDict,
NamedTuple, msgspec.Struct, etc.).
from litestar import Litestar, websocket_listener
@websocket_listener("/")
async def handler(data: str) -> str:
return data
app = Litestar([handler])
from litestar import Litestar, websocket_listener
@websocket_listener("/")
async def handler(data: str) -> bytes:
return data.encode("utf-8")
app = Litestar([handler])
from typing import Dict
from litestar import Litestar, websocket_listener
@websocket_listener("/")
async def handler(data: str) -> Dict[str, str]:
return {"message": data}
app = Litestar([handler])
from litestar import Litestar, websocket_listener
@websocket_listener("/")
async def handler(data: str) -> dict[str, str]:
return {"message": data}
app = Litestar([handler])
from dataclasses import dataclass
from datetime import datetime
from litestar import Litestar, websocket_listener
@dataclass
class Message:
content: str
timestamp: float
@websocket_listener("/")
async def handler(data: str) -> Message:
return Message(content=data, timestamp=datetime.now().timestamp())
app = Litestar([handler])
Transport modes#
WebSockets have two transport modes: Text and binary. These can be specified individually for receiving and sending data.
Note
It may seem intuitive that text and binary should map to str and
bytes respectively, but this is not the case. Listeners can receive and
send data in any format, independently of the mode. The mode only affects how
data is encoded during transport (i.e. on the protocol level). In most cases the
default mode - text - is all that’s needed. Binary transport is usually employed
when sending binary blobs that don’t have a meaningful string representation, such
as images.
Setting the receive mode#
text is the default mode and is appropriate for most messages, including
structured data such as JSON.
from litestar import Litestar, websocket_listener
@websocket_listener("/", receive_mode="text")
async def handler(data: str) -> str:
return data
app = Litestar([handler])
from litestar import Litestar, websocket_listener
@websocket_listener("/", receive_mode="binary")
async def handler(data: str) -> str:
return data
app = Litestar([handler])
Important
Once configured with a mode, a listener will only listen to socket events of the
appropriate type. This means if a listener is configured to use binary mode,
it will not respond to websocket events sending data in the text channel
Setting the send mode#
text is the default mode and is appropriate for most messages, including
structured data such as JSON.
from litestar import Litestar, websocket_listener
@websocket_listener("/", send_mode="text")
async def handler(data: str) -> str:
return data
app = Litestar([handler])
from litestar import Litestar, websocket_listener
@websocket_listener("/", send_mode="binary")
async def handler(data: str) -> str:
return data
app = Litestar([handler])
Dependency injection#
Dependency Injection is available as well and generally works the same as with regular route handlers:
from litestar import Litestar, websocket_listener
from litestar.di import Provide
def some_dependency() -> str:
return "hello"
@websocket_listener("/", dependencies={"some": Provide(some_dependency)})
async def handler(data: str, some: str) -> str:
return data + some
app = Litestar([handler])
Important
Injected dependencies work on the level of the underlying route handler. This means they won’t be re-evaluated every time the listener function is called.
The following example makes use of yield dependencies and
the fact that dependencies are only evaluated once for every connection; The step after
the yield will only be executed after the connection has been closed.
from typing import TypedDict
from litestar import Litestar, websocket_listener
from litestar.datastructures import State
from litestar.di import Provide
class Message(TypedDict):
message: str
client_count: int
def socket_client_count(state: State) -> int:
if not hasattr(state, "count"):
state.count = 0
state.count += 1
yield state.count
state.count -= 1
@websocket_listener("/", dependencies={"client_count": Provide(socket_client_count)})
async def handler(data: str, client_count: int) -> Message:
return Message(message=data, client_count=client_count)
app = Litestar([handler])
Interacting with the WebSocket directly#
Sometimes access to the socket instance is needed, in which case the
WebSocket instance can be injected into the handler
function via the socket argument:
from litestar import Litestar, WebSocket, websocket_listener
@websocket_listener("/")
async def handler(data: str, socket: WebSocket) -> str:
if data == "goodbye":
await socket.close()
return data
app = Litestar([handler])
Important
Since websockets are inherently asynchronous, to interact with the asynchronous
methods on WebSocket, the handler function needs
to be asynchronous.
Customising connection acceptance#
By default, Litestar will accept all incoming connections by awaiting WebSocket.accept() without arguments.
This behavior can be customized by passing a custom connection_accept_handler function. Litestar will await this
function to accept the connection.
from litestar import Litestar, WebSocket, websocket_listener
async def accept_connection(socket: WebSocket) -> None:
await socket.accept(headers={"Cookie": "custom-cookie"})
@websocket_listener("/", connection_accept_handler=accept_connection)
def handler(data: str) -> str:
return data
app = Litestar([handler])
Class based WebSocket handling#
In addition to using a simple function as in the examples above, a class based approach
is made possible by extending the
WebSocketListener. This provides
convenient access to socket events such as connect and disconnect, and can be used to
encapsulate more complex logic.
from litestar import Litestar, WebSocket
from litestar.handlers import WebsocketListener
class Handler(WebsocketListener):
path = "/"
def on_accept(self, socket: WebSocket) -> None:
print("Connection accepted")
def on_disconnect(self, socket: WebSocket) -> None:
print("Connection closed")
def on_receive(self, data: str) -> str:
return data
app = Litestar([Handler])
from litestar import Litestar, WebSocket
from litestar.handlers import WebsocketListener
class Handler(WebsocketListener):
path = "/"
async def on_accept(self, socket: WebSocket) -> None:
print("Connection accepted")
async def on_disconnect(self, socket: WebSocket) -> None:
print("Connection closed")
async def on_receive(self, data: str) -> str:
return data
app = Litestar([Handler])
Custom WebSocket#
New in version 2.7.0.
Litestar supports custom websocket_class instances, which can be used to further configure the default WebSocket.
The example below illustrates how to implement custom websocket class for the whole application.
Example of a custom websocket at the application level
from __future__ import annotations
from litestar import Litestar, WebSocket, websocket_listener
from litestar.types.asgi_types import WebSocketMode
class CustomWebSocket(WebSocket):
async def receive_data(self, mode: WebSocketMode) -> str | bytes:
"""Return fixed response for every websocket message."""
await super().receive_data(mode=mode)
return "Fixed response"
@websocket_listener("/")
async def handler(data: str) -> str:
return data
app = Litestar([handler], websocket_class=CustomWebSocket)
Layered architecture
WebSocket classes are part of Litestar’s layered architecture, which means you can set a websocket class on every layer of the application. If you have set a webscoket class on multiple layers, the layer closest to the route handler will take precedence.
You can read more about this in the Layered architecture section