mtapi/PyMtApi5/mtrpcclient.py

import logging
from enum import IntEnum
from threading import Condition, Lock, Thread

import websockets
from websockets.sync.client import connect as ws_connect


class MtNotification(IntEnum):
    ClientReady = 0


class MtMessageType(IntEnum):
    Command = 0
    Response = 1
    Event = 2
    ExpertList = 3
    ExpertAdded = 4
    ExpertRemoved = 5
    Notification = 6


class CommandTask:
    def __init__(self):
        self.locker = Lock()
        self.waiter = Condition()
        self.response = None

    def wait_response(self, time):
        with self.waiter:
            self.waiter.wait(time)
        with self.locker:
            return self.response

    def set_response(self, response):
        with self.locker:
            self.response = response
        with self.waiter:
            self.waiter.notify()


class MtRpcClient:
    def __init__(self, callback=None):
        self.__logger = logging.getLogger(__name__)
        self.__callback = callback
        self.__notification_tasks = dict()
        self.__tasks = dict()
        self.__next_command_id = 0
        self.__lock = Lock()

    def connect(self, url):
        self.__logger.debug(f"connecting to {url}")
        self.__ws = ws_connect(url)
        self.__receive_thread = Thread(target=self.__receive_messages_thread)
        self.__receive_thread.start()

    def disconnect(self):
        self.__ws.close()
        self.__receive_thread.join()
        self.__logger.debug("disconnected")

    def request_expert_list(self):
        task = CommandTask()
        with self.__lock:
            self.__notification_tasks[MtNotification.ClientReady] = task
        self.__ws.send(self.__create_notification(MtNotification.ClientReady))
        response = task.wait_response(10)
        with self.__lock:
            self.__notification_tasks.pop(MtNotification.ClientReady)
        return response

    def send_command(self, expert_handle, command_type, payload=None):
        command_id = self.__next_command_id
        self.__next_command_id += 1
        task = CommandTask()
        with self.__lock:
            self.__tasks[command_id] = task
        self.__ws.send(self.__create_mt_command(expert_handle, command_id, command_type, payload))
        response = task.wait_response(10)
        with self.__lock:
            self.__tasks.pop(command_id)
        return response

    # Private methods

    def __process_message(self, message):
        self.__logger.debug(f"process_message: {message}")
        pieces = message.split(";", 1)
        if len(pieces) != 2 or not pieces[0] or not pieces[1]:
            self.__logger.warning("process_message: Invalid message format")
            return
        message_type = MtMessageType(int(pieces[0]))
        if message_type == MtMessageType.ExpertList:
            self.__process_expert_list(pieces[1])
        elif message_type == MtMessageType.Event:
            self.__process_event(pieces[1])
        elif message_type == MtMessageType.Response:
            self.__process_response(pieces[1])
        elif message_type == MtMessageType.ExpertAdded:
            self.__process_expert_added(pieces[1])
        elif message_type == MtMessageType.ExpertRemoved:
            self.__process_expert_removed(pieces[1])
        else:
            self.__logger.warning(f"received unknown message type: {message_type}")

    def __process_expert_list(self, payload):
        pieces = payload.split(",")
        experts = list()
        for p in pieces:
            experts.append(int(p))
        with self.__lock:
            task = self.__notification_tasks.get(MtNotification.ClientReady)
            if task is not None:
                task.set_response(experts)

    def __process_event(self, payload):
        pieces = payload.split(";", 2)
        if len(pieces) != 3 or not pieces[0] or not pieces[1] or not pieces[2]:
            self.__logger.warning("process_event: Invalid message format")
            return
        if self.__callback is not None:
            self.__callback.mt_rpc_on_event(int(pieces[0]), int(pieces[1]), pieces[2])

    def __process_response(self, payload):
        pieces = payload.split(";", 2)
        if len(pieces) != 3 or not pieces[0] or not pieces[1] or not pieces[2]:
            self.__logger.warning("process_response: Invalid message format")
            return
        command_id = int(pieces[1])
        with self.__lock:
            task = self.__tasks.get(command_id)
            if task is not None:
                task.set_response(pieces[2])

    def __process_expert_added(self, payload):
        if self.__callback is not None:
            self.__callback.mt_rpc_on_expert_added(int(payload))

    def __process_expert_removed(self, payload):
        if self.__callback is not None:
            self.__callback.mt_rpc_on_expert_removed(int(payload))

    def __receive_messages_thread(self):
        self.__logger.debug("started receive_messages thread")
        while True:
            try:
                message = self.__ws.recv()
                self.__process_message(message)
            except websockets.exceptions.ConnectionClosed:
                self.__logger.info("Connection closed")
                if self.__callback is not None:
                    self.__callback.mt_rcp_on_disconnect()
                break
            except Exception as e:
                self.__logger.error(e)
                if self.__callback is not None:
                    self.__callback.mt_rpc_on_connection_failed(str(e))
                break
        self.__logger.debug("function receive_messages finished")

    def __create_notification(self, notification_type):
        return f"{int(MtMessageType.Notification)};{notification_type}"

    def __create_mt_command(self, expert_handle, command_id, command_type, payload):
        if payload is None:
            return f"{MtMessageType.Command};{expert_handle};{command_id};{command_type}"
        return f"{MtMessageType.Command};{expert_handle};{command_id};{command_type};{payload}"