Add event callback to async connection

samsungtvws/async_connection.py

@@ -22,11 +22,14 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import json
 import logging
 import ssl
+from typing import Any, Awaitable, Callable
 
 from websockets.client import connect
+from websockets.exceptions import ConnectionClosed
 
 from . import connection, exceptions, helper
 from .command import SamsungTVCommand
@@ -57,9 +60,30 @@ async def open(self):
 
         return connection
 
+    async def start_listening(
+        self, callback: Callable[[str, Any], Awaitable[None]]
+    ) -> None:
+        """Open, and start listening."""
+        if self.connection is None:
+            self.connection = await self.open()
+
+        self._recv_loop = asyncio.create_task(self._do_start_listening(callback))
+
+    async def _do_start_listening(
+        self, callback: Callable[[str, Any], Awaitable[None]]
+    ) -> None:
+        """Do start listening."""
+        with contextlib.suppress(ConnectionClosed):
+            while True:
+                data = await self.connection.recv()
+                response = helper.process_api_response(data)
+                await callback(response.get("event", "*"), response)
+
     async def close(self):
         if self.connection:
             await self.connection.close()
+            if self._recv_loop:
+                await self._recv_loop
 
         self.connection = None
         _LOGGING.debug("Connection closed.")

samsungtvws/connection.py

@@ -19,10 +19,12 @@
     Boston, MA  02110-1335  USA
 
 """
+import contextlib
 import json
 import logging
 import ssl
 import time
+from typing import Any, Callable
 
 import websocket
 
@@ -60,6 +62,7 @@ def __init__(
         self.name = name
         self.endpoint = endpoint
         self.connection = None
+        self._recv_loop = None
 
     def _is_ssl_connection(self):
         return self.port == 8002
@@ -146,6 +149,19 @@ def open(self):
 
         return connection
 
+    def start_listening(self, callback: Callable[[str, Any], None]) -> None:
+        """Start listening (use with asyncio.create_task())."""
+        if self.connection is None:
+            self.connection = self.open()
+
+    def _start_listening(self, callback: Callable[[str, Any], None]) -> None:
+        """Start listening (use with asyncio.create_task())."""
+        with contextlib.suppress(websocket.WebSocketException):
+            while data := self.connection.recv():
+                response = helper.process_api_response(data)
+                callback(response.get("event", "*"), response)
+                time.sleep(1)
+
     def close(self):
         if self.connection:
             self.connection.close()