From 5aed29b20af32085da8993db0dd76e54424c349e Mon Sep 17 00:00:00 2001 From: Sergey Sayamov Date: Thu, 7 Mar 2024 23:11:44 +0300 Subject: [PATCH] Allow user to redefine websocket.accept in BaseWSEndpoint --- starlette_web/common/ws/base_endpoint.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/starlette_web/common/ws/base_endpoint.py b/starlette_web/common/ws/base_endpoint.py index 945dbe2..7bc894c 100644 --- a/starlette_web/common/ws/base_endpoint.py +++ b/starlette_web/common/ws/base_endpoint.py @@ -61,10 +61,13 @@ async def on_connect(self, websocket: WebSocket) -> None: del websocket.state.db_session if permitted: - await websocket.accept() + await self.accept(websocket) else: raise WebSocketDisconnect(code=3000, reason=reason) + async def accept(self, websocket: WebSocket) -> None: + await websocket.accept() + async def on_receive(self, websocket: WebSocket, data: Any) -> None: cleaned_data = self._validate(data) task_id = get_random_string(50)