From 08553f84b9319ada8bd5cc3c90b8457fdc87ca72 Mon Sep 17 00:00:00 2001 From: Fabio Manganiello Date: Wed, 15 Feb 2023 22:23:15 +0100 Subject: [PATCH] Added `timeout` parameter to `websocket.send`. --- platypush/plugins/websocket/__init__.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/platypush/plugins/websocket/__init__.py b/platypush/plugins/websocket/__init__.py index 5744c075d8..0f4eb49f5c 100644 --- a/platypush/plugins/websocket/__init__.py +++ b/platypush/plugins/websocket/__init__.py @@ -50,6 +50,7 @@ class WebsocketPlugin(AsyncRunnablePlugin): ssl_cafile=None, ssl_capath=None, wait_response=False, + timeout=None, ): """ Sends a message to a websocket. @@ -66,6 +67,9 @@ class WebsocketPlugin(AsyncRunnablePlugin): required by the SSL configuration (default: None) :param wait_response: Set to True if you expect a response to the delivered message. + :param timeout: If ``wait_response=True``, then ``timeout`` establishes + how long we should wait for a response before returning (default: + no timeout). :return: The received response if ``wait_response`` is set to True, otherwise nothing. """ @@ -87,10 +91,13 @@ class WebsocketPlugin(AsyncRunnablePlugin): try: await ws.send(str(msg)) except ConnectionClosed as err: - self.logger.warning('Error on websocket %s: %s', url, err) + self.logger.warning( + 'Connection error to websocket %s: %s', url, err + ) if wait_response: - messages = await self._recv(ws, num_messages=1) + messages = await self._recv(ws, num_messages=1, timeout=timeout) + if messages: return self._parse_msg(messages[0]) @@ -145,7 +152,7 @@ class WebsocketPlugin(AsyncRunnablePlugin): return self.loop.call_soon_threadsafe(recv) - async def _recv(self, ws, timeout=0, num_messages=0): + async def _recv(self, ws, timeout: Optional[float] = 0, num_messages=0): messages = [] time_start = time.time() time_end = time_start + timeout if timeout else 0