feat: replace asyncio websocket with a more easy to use threading one (#87)

Author: LouisBrunner

Pipfile

@@ -6,7 +6,7 @@ name = "pypi"
 [packages]
 types-requests = "*"
 requests = "*"
-websockets = "*"
+"websocket-client" = "*"
 zeroconf = "*"
 types-dataclasses = "*"
 dataclasses = "*"

Pipfile.lock

@@ -3,6 +3,7 @@
 from .helpers import (strip_ip)
 from .eva_http_client import EvaHTTPClient
 from .eva_locker import EvaWithLocker
+from .eva_ws import Websocket
 
 
 class Eva:
@@ -29,6 +30,15 @@ def set_request_timeout(self, request_timeout):
     __TEACH_RENEW_PERIOD = 3
 
 
+    # --------------------------------------------- Websocket -----------------------------------------------
+    def websocket(self):
+        # Ensure we have a session token
+        # might use the result in the future to give the initial state to consumers
+        self.__http_client.data_snapshot()
+        host_uri = f'ws://{self.__http_client.host_ip}/api/v1/data/stream'
+        subprotocols = [f'SessionToken_{self.__http_client.session_token}', "object"]
+        return Websocket(host_uri, subprotocols, timeout=self.__http_client.request_timeout)
+
     # --------------------------------------------- Lock Holder ---------------------------------------------
     def __enter__(self):
         self.__eva_locker.__enter__()

evasdk/Eva.py

@@ -69,13 +69,12 @@
 
 from .Eva import Eva
 from .eva_http_client import EvaHTTPClient
-from .eva_ws import ws_connect
 from .robot_state import RobotState
 from .helpers import strip_ip
 from .eva_errors import (
     EvaError,
     EvaValidationError, EvaAuthError, EvaAutoRenewError,
-    EvaAdminError, EvaServerError)
+    EvaAdminError, EvaServerError, EvaWebsocketError)
 from .version import __version__
 from .EvaDiscoverer import (
     DiscoverCallback, DiscoveredEva,

evasdk/__init__.py

@@ -30,6 +30,10 @@ class EvaAutoRenewError(EvaError, Exception):
     """Error thrown when automatic session renewal fails but not the original request"""
 
 
+class EvaWebsocketError(EvaError, Exception):
+    """Error thrown when an issue related to the websocket stream happens"""
+
+
 def eva_error(label, r=None):
     if r is not None:
         __handle_http_error(label, r)

evasdk/eva_errors.py

@@ -1,17 +1,89 @@
-import websockets
+import websocket  # type: ignore
+import logging
+import json
+from threading import Thread, Condition
+from .eva_errors import EvaWebsocketError
+from .observer import Subject
 
 
-async def ws_connect(host_ip, session_token):
-    """
-    Connect is an async function that returns a connected Eva websocket
+logger = logging.getLogger(__name__)
+
 
-    Connect needs to be run from a asyncio event_loop and retuns a
-    websockets.Websocket object. Using this you can manually call .recv()
-    and .send(), make sure to also .close() when you are finished to clean
-    up the websocket connection.
+class Websocket:
     """
-    host_uri = 'ws://{}/api/v1/data/stream'.format(host_ip)
-    subprotocols = ['SessionToken_{}'.format(session_token), "object"]
+    This class creates a context which runs a thread to monitor a websocket in the background.
+    It follows the Observer pattern, which you can use to listen to specific event or `all` to get all of them.
+    Note: notifications will be sent from a different thread so you will need to use a mutex or other synchronization mechanims.
+    """
+
+    def __init__(self, url, protocols, timeout=5):
+        self.__thread = Thread(target=self.__run)
+        self.__url = url
+        self.__protocols = protocols
+        self.__app = None
+        self.__ws = None
+        self.__cond = Condition()
+        self.__timeout = timeout
+        self.__subject = Subject()
+
+    def __enter__(self):
+        self.__thread.start()
+        with self.__cond:
+            connected = self.__cond.wait(timeout=self.__timeout)
+            if not connected:
+                raise EvaWebsocketError('could not connect to Eva\'s data stream')
+        return self
+
+    def __on_open(self, ws):
+        logger.debug('ws_on_open')
+        with self.__cond:
+            self.__ws = ws
+            self.__cond.notify_all()
+
+    def __on_message(self, ws, raw):
+        msg = json.loads(raw)
+        if 'type' not in msg or msg['type'] != 'state_change':
+            logger.debug('ws_on_message', msg)
+        self.__subject.notify('all', msg)
+        if 'type' in msg:
+            self.__subject.notify(msg['type'], msg)
+        else:
+            logger.debug('missing message type', msg)
+
+    def __on_error(self, ws, error):
+        logger.debug('ws_on_error', error)
+        with self.__cond:
+            self.__ws = None
+
+    def __on_close(self, ws, close_status_code, close_msg):
+        logger.debug('ws_on_close', close_status_code, close_msg)
+        with self.__cond:
+            self.__ws = None
+
+    def __exit__(self, type, value, traceback):
+        if self.__app is not None:
+            self.__app.keep_running = False
+        self.__thread.join()
+
+    def __run(self):
+        self.__app = websocket.WebSocketApp(
+            self.__url,
+            subprotocols=self.__protocols,
+            on_open=self.__on_open,
+            on_message=self.__on_message,
+            on_error=self.__on_error,
+            on_close=self.__on_close,
+        )
+        self.__app.run_forever()
+
+    def _send_raw(self, data):
+        with self.__cond:
+            if self.__ws is None:
+                raise EvaWebsocketError('Eva\'s data stream has disconnected')
+            self.__ws.send(json.dumps(data))
+
+    def register(self, event, callback):
+        self.__subject.register(event, callback)
 
-    ws = await websockets.client.connect(host_uri, subprotocols=subprotocols)
-    return ws
+    def deregister(self, event, callback):
+        self.__subject.deregister(event, callback)

evasdk/eva_ws.py

@@ -0,0 +1,21 @@
+class Subject:
+    def __init__(self):
+        self.events = {}
+
+    def register(self, event, callback):
+        if event not in self.events:
+            self.events[event] = set()
+        self.events[event].add(callback)
+
+    def deregister(self, event, callback):
+        if event not in self.events:
+            return
+        if callback not in self.events[event]:
+            return
+        self.events[event].remove(callback)
+
+    def notify(self, event, *args):
+        if event not in self.events:
+            return
+        for callback in self.events[event]:
+            callback(*args)

evasdk/observer.py

@@ -1,9 +1,6 @@
 #!/usr/bin/env python3
 
-import evasdk
-import json
-import asyncio
-import time
+from evasdk import Eva
 
 # This example shows usage of the eva_ws and eva_http modules, used for direct control
 # using the network interfaces. eva_http also contains some helper functions not
@@ -12,33 +9,16 @@
 host_ip = input("Please enter a Eva IP: ")
 token = input("Please enter a valid Eva token: ")
 
-print('ip: [{}], token: [{}]\n'.format(host_ip, token))
+print(f'ip: [{host_ip}], token: [{token}]\n')
 
-http_client = evasdk.EvaHTTPClient(host_ip, token)
+eva = Eva(host_ip, token)
 
-# The session token will be valid for 30 minutes, you'll need to renew the session
-# if you want the websocket connection to continue after that point.
-session_token = http_client.auth_create_session()
+users = eva.users_get()
+print(f'Eva at {host_ip} users: {users}\n')
 
-users = http_client.users_get()
-print('Eva at {} users: {}\n'.format(host_ip, users))
+joint_angles = eva.data_servo_positions()
+print(f'Eva current joint angles: {joint_angles}')
 
-joint_angles = http_client.data_servo_positions()
-print('Eva current joint angles: {}'.format(joint_angles))
-
-
-async def eva_ws_example(host_ip, session_token):
-    websocket = await evasdk.ws_connect(host_ip, session_token)
-
-    msg_count = 0
-    time_since_msg = time.time()
-    while True:
-        ws_msg_json = await websocket.recv()
-        print('WS msg delta T: {}, number: {}'.format(time.time() - time_since_msg, msg_count))
-        msg_count += 1
-        time_since_msg = time.time()
-
-        ws_msg = json.loads(ws_msg_json)
-        print(ws_msg)
-
-asyncio.get_event_loop().run_until_complete(eva_ws_example(host_ip, session_token))
+with eva.websocket() as ws:
+    ws.register('state_change', print)
+    input('press return when you want to stop\n')

examples/http_ws_example.py

@@ -18,7 +18,7 @@
     long_description_content_type="text/markdown",
     install_requires=[
         'requests',
-        'websockets',
+        'websocket-client',
         'zeroconf',
         'dataclasses',
         # TODO: too big, install it manually if you want it