Coverage for src/lexigram/web/websocket/gateway.py: 28%
104 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 04:37 +0800
1"""WebSocket Gateway — declarative event-driven WebSocket handling.
3Introduces a higher-level ``WebSocketGateway`` base class and companion
4decorators that replace manual ``if/elif`` dispatch inside ``on_message``
5with a registry-based, type-safe event routing system.
7Example::
9 from lexigram.web.websocket.gateway import (
10 WebSocketGateway,
11 on_connect,
12 subscribe_message,
13 websocket_gateway,
14 )
16 @websocket_gateway("/ws/chat")
17 class ChatGateway(WebSocketGateway):
19 def __init__(self, rooms: RoomManager) -> None:
20 super().__init__()
21 self.rooms = rooms # DI-resolved
23 @subscribe_message("chat.send")
24 async def handle_send(self, websocket: WebSocket, payload: dict) -> None:
25 await self.rooms.broadcast(payload)
27 @subscribe_message("room.join")
28 async def handle_join(self, websocket: WebSocket, payload: dict) -> None:
29 await self.rooms.join(payload.get("room"), websocket)
31 @on_connect
32 async def connected(self, websocket: WebSocket) -> None:
33 await websocket.accept()
34"""
36from __future__ import annotations
38from collections.abc import Callable
39from typing import TYPE_CHECKING, Any, TypeVar
41from lexigram.logging import get_logger
42from lexigram.web.websocket.handler import AbstractWebSocketHandler
44if TYPE_CHECKING:
45 from lexigram.web.transport.websockets import WebSocket
47logger = get_logger(__name__)
49F = TypeVar("F", bound=Callable[..., Any])
51# Sentinel attribute names written onto decorated methods
52_SUBSCRIBE_MESSAGE_ATTR = "__ws_subscribe_message__"
53_ON_CONNECT_ATTR = "__ws_on_connect__"
54_ON_DISCONNECT_ATTR = "__ws_on_disconnect__"
55_ON_ERROR_ATTR = "__ws_on_error__"
58def subscribe_message(event_type: str) -> Callable[[F], F]:
59 """Mark a gateway method as the handler for a specific message type.
61 The method will be called automatically when an incoming message's
62 ``"type"`` field (or top-level key) matches ``event_type``.
64 Args:
65 event_type: The message type string to match (e.g. ``"chat.send"``).
67 Example::
69 @subscribe_message("chat.send")
70 async def handle_send(self, websocket: WebSocket, payload: dict) -> None:
71 ...
72 """
74 def decorator(fn: F) -> F:
75 existing: list[str] = getattr(fn, _SUBSCRIBE_MESSAGE_ATTR, [])
76 fn.__ws_subscribe_message__ = [*existing, event_type] # type: ignore[attr-defined]
77 return fn
79 return decorator
82def on_connect(fn: F) -> F:
83 """Mark a gateway method as the connection handler.
85 The decorated method replaces ``on_connect`` in :class:`AbstractWebSocketHandler`.
86 Only one ``@on_connect`` method per gateway is allowed.
88 Args:
89 fn: The async method to call on connection.
90 """
91 fn.__ws_on_connect__ = True # type: ignore[attr-defined]
92 return fn
95def on_disconnect(fn: F) -> F:
96 """Mark a gateway method as the disconnection handler.
98 Args:
99 fn: The async method to call upon disconnection.
100 """
101 fn.__ws_on_disconnect__ = True # type: ignore[attr-defined]
102 return fn
105def on_error(fn: F) -> F:
106 """Mark a gateway method as the error handler.
108 Args:
109 fn: The async method to call when an error occurs.
110 """
111 fn.__ws_on_error__ = True # type: ignore[attr-defined]
112 return fn
115class WebSocketGateway(AbstractWebSocketHandler):
116 """High-level declarative WebSocket gateway.
118 Extend this class and use :func:`subscribe_message`, :func:`on_connect`,
119 and :func:`on_disconnect` decorators to define event handlers. The
120 gateway automatically dispatches incoming messages to the correct method
121 based on the ``"type"`` field of the received JSON payload.
123 If no specific handler matches, the default ``on_unhandled_message``
124 method is called (no-op by default; override to customise).
126 DI works naturally — declare constructor parameters, and the container
127 will inject dependencies when the gateway is resolved::
129 @websocket_gateway("/ws/chat")
130 class ChatGateway(WebSocketGateway):
132 def __init__(self, rooms: RoomManager) -> None:
133 super().__init__()
134 self.rooms = rooms
136 @subscribe_message("chat.send")
137 async def handle_send(self, ws: WebSocket, payload: dict) -> None:
138 await self.rooms.broadcast(payload)
139 """
141 _message_registry: dict[str, Callable[..., Any]] = {}
142 _connect_handler: Callable[..., Any] | None = None
143 _disconnect_handler: Callable[..., Any] | None = None
144 _error_handler: Callable[..., Any] | None = None
146 def __init_subclass__(cls, **kwargs: Any) -> None:
147 super().__init_subclass__(**kwargs)
148 # Build the message dispatch registry for this subclass
149 registry: dict[str, Callable[..., Any]] = {}
150 connect_handler: Callable[..., Any] | None = None
151 disconnect_handler: Callable[..., Any] | None = None
152 error_handler: Callable[..., Any] | None = None
154 for name in dir(cls):
155 try:
156 method = getattr(cls, name)
157 except AttributeError:
158 continue
160 # Message subscriptions — a method can handle multiple event types
161 subscriptions: list[str] = getattr(method, _SUBSCRIBE_MESSAGE_ATTR, [])
162 for event_type in subscriptions:
163 if event_type in registry:
164 logger.warning(
165 "Duplicate @subscribe_message handler for type %r in %s",
166 event_type,
167 cls.__name__,
168 )
169 registry[event_type] = method
171 if getattr(method, _ON_CONNECT_ATTR, False):
172 connect_handler = method
173 if getattr(method, _ON_DISCONNECT_ATTR, False):
174 disconnect_handler = method
175 if getattr(method, _ON_ERROR_ATTR, False):
176 error_handler = method
178 cls._message_registry = registry
179 cls._connect_handler = connect_handler
180 cls._disconnect_handler = disconnect_handler
181 cls._error_handler = error_handler
183 def __init__(self) -> None:
184 super().__init__()
185 # Ensure per-instance registry copy so subclass attributes shadow safely
186 self._message_registry = dict(type(self)._message_registry)
188 async def on_connect(self, websocket: WebSocket) -> None:
189 """Routes to the @on_connect decorated method, or accepts by default."""
190 connect = type(self)._connect_handler
191 if connect is not None:
192 await connect(self, websocket)
193 else:
194 # Safe default — accept and continue
195 await websocket.accept()
197 async def on_disconnect(self, websocket: WebSocket) -> None:
198 """Routes to the @on_disconnect decorated method, or no-ops."""
199 disconnect = type(self)._disconnect_handler
200 if disconnect is not None:
201 await disconnect(self, websocket)
203 async def on_error(self, websocket: WebSocket, error: Exception) -> None:
204 """Routes to the @on_error decorated method, or uses default handling."""
205 error_h = type(self)._error_handler
206 if error_h is not None:
207 await error_h(self, websocket, error)
208 else:
209 await super().on_error(websocket, error)
211 async def on_message(self, websocket: WebSocket, message: dict[str, Any]) -> None:
212 """Dispatch incoming message to the correct @subscribe_message handler.
214 Looks up the ``"type"`` field of the message in the registry.
215 If no match is found, ``on_unhandled_message`` is called.
217 Args:
218 websocket: The WebSocket connection.
219 message: The received JSON payload (already parsed).
220 """
221 event_type = message.get("type") or message.get("event")
222 if event_type is None:
223 await self.on_unhandled_message(websocket, message)
224 return
226 handler = self._message_registry.get(str(event_type))
227 if handler is None:
228 logger.debug("No handler for WebSocket event %r", event_type)
229 await self.on_unhandled_message(websocket, message)
230 return
232 # Strip the "type" key from the payload before passing to the handler
233 payload = {k: v for k, v in message.items() if k not in ("type", "event")}
234 await handler(self, websocket, payload)
236 async def on_unhandled_message(
237 self,
238 websocket: WebSocket,
239 message: dict[str, Any],
240 ) -> None:
241 """Called when no handler matches the incoming message type.
243 Override to customise the behaviour. Default sends an error response.
245 Args:
246 websocket: The WebSocket connection.
247 message: The unhandled message.
248 """
249 event_type = message.get("type") or message.get("event") or "unknown"
250 await self.send_error(websocket, f"Unhandled event: {event_type}")
253def websocket_gateway(
254 path: str,
255 *,
256 ping_interval: int | None = None,
257 ping_timeout: int | None = None,
258 max_connections_per_user: int | None = None,
259) -> Callable[[type], type]:
260 """Mark a class as a WebSocket gateway endpoint.
262 This is the high-level companion to :func:`lexigram.web.websocket.decorators.websocket_handler`.
263 Use it together with :class:`WebSocketGateway`.
265 Args:
266 path: URL path for the WebSocket endpoint (may include path params).
267 ping_interval: Override default keepalive ping interval (seconds).
268 ping_timeout: Override default pong timeout (seconds).
269 max_connections_per_user: Limit concurrent connections per user.
271 Example::
273 @websocket_gateway("/ws/notifications")
274 class NotificationGateway(WebSocketGateway):
275 ...
276 """
278 def decorator(cls: type) -> type:
279 if ping_interval is not None:
280 cls.ping_interval = ping_interval # type: ignore[attr-defined]
281 if ping_timeout is not None:
282 cls.ping_timeout = ping_timeout # type: ignore[attr-defined]
283 if max_connections_per_user is not None:
284 cls.max_connections_per_user = max_connections_per_user # type: ignore[attr-defined]
286 cls._ws_path = path # type: ignore[attr-defined]
287 cls._ws_metadata = { # type: ignore[attr-defined]
288 "path": path,
289 "ping_interval": ping_interval,
290 "ping_timeout": ping_timeout,
291 "max_connections_per_user": max_connections_per_user,
292 }
293 cls._is_websocket_handler = True # type: ignore[attr-defined]
294 cls._is_websocket_gateway = True # type: ignore[attr-defined]
295 return cls
297 return decorator
300__all__ = [
301 "WebSocketGateway",
302 "on_connect",
303 "on_disconnect",
304 "on_error",
305 "subscribe_message",
306 "websocket_gateway",
307]