pqAutomationApp/UniTAP/dev/device.py

import copy
import weakref
from typing import TypeVar, Type, List

from .dev_5xx_roles import UCD500
from .dev_4xx_roles import UCD400, UCD422, UCD424
from .dev_3xx_roles import UCD240, UCD340, UCD323, UCD301
from .modules import *

from UniTAP.libs.lib_tsi.tsi_io import DeviceIO
from UniTAP.utils import tsi_logging as logging

RoleType = TypeVar("RoleType",
                   UCD500.DPSourceUSBCSink,
                   UCD500.DPSourceDPSink,
                   UCD500.USBCSourceUSBCSink,
                   UCD500.USBCSourceDPSink,

                   UCD424.USBCSourceUSBCSink,
                   UCD422.HDMISourceHDMISink,
                   UCD400.DPSourceDPSink,

                   UCD340.USBCSink,
                   UCD340.USBCSource,

                   UCD323.DPSink,
                   UCD323.DPSource,
                   UCD323.HDMISink,
                   UCD323.HDMISource,

                   UCD301.DPSink,
                   UCD301.HDMISink,

                   UCD240.USBCSink,
                   UCD240.USBCSource)

MODEL_TO_CLASS = {
    "UCD-500": UCD500,
    "UCD-424": UCD424,
    "UCD-422": UCD422,
    "UCD-400": UCD400,
    "UCD-340": UCD340,
    "UCD-323": UCD323,
    "UCD-301": UCD301,
    "UCD-240": UCD240
}


class NotExistingRoleSelected(Exception):
    """
    Redefinition of base exception.
    Define error of 'not existing role selected'.
    """
    pass


class DeviceRoleIsAlreadySelected(Exception):
    """
    Redefinition of base exception.
    Define error of 'device role already selected'.
    """
    pass


class TSIDevice:
    """
    Class `TSIDevice` ... Allows:
    - Select role of the device `select_role`.
    - Get and redefine OPF (Operator feedback dialog) handler `opf_handler`.
    """

    def __init__(self, io: DeviceIO, name: str, serial_number: str, roles: List[RoleType]):
        self.__io = io
        self.__io_proxy = weakref.proxy(self.__io)
        self.__memory_manager = MemoryManager(self.__io_proxy)
        self.__capturer = Capturer(self.__io_proxy)
        self.__opf_handler = OperatorFeedbackHandler(self.__io_proxy)
        self.__terminal = Terminal(self.__io_proxy)
        self.__name = name
        self.__serial_number = serial_number
        self.__role_list = roles
        logging.info(f"[UniTAP] TSIDevice.init: {self}")

    def __del__(self):
        del self.__io
        logging.info(f"[UniTAP] TSIDevice.del: {self}")

    def __str__(self):
        return f"{self.__name} [{self.__serial_number}]"

    def __define_internal_handler(self, role, role_type):
        if role_type in [UCD500.DPSourceUSBCSink, UCD500.DPSourceDPSink, UCD500.USBCSourceUSBCSink,
                         UCD500.USBCSourceDPSink, UCD424.USBCSourceUSBCSink, UCD400.DPSourceDPSink]:
            self.opf_handler._set_roles(role.dptx, role.dprx)
        elif role_type in [UCD422.HDMISourceHDMISink]:
            self.opf_handler._set_roles(role.hdtx, role.hdrx)

    def select_role(self, role_type: Type[RoleType]) -> RoleType:
        """
        Function allows selecting role of the chosen device.
        For example: `UniTAP.dev.UCD500.DPSourceDPSink` - will be opened UCD-500 in role DisplayPort Source and Sink.

        Args:
            role_type (RoleType) - one of the possible roles.

        Returns:
            object of `RoleType` type. Selected role.
        """
        # TODO: Remove this limitation. Caused by TSI lib problem.
        if self.__io.get_role() is not None:
            raise DeviceRoleIsAlreadySelected(f"Device role is already selected.")
        if role_type not in self.__role_list:
            raise NotExistingRoleSelected(f"Role: {role_type} is not available on this device.")
        self.__io.select_role(self.__role_list.index(role_type))
        self.__io.set_role(role_type(self.__io_proxy,
                                     weakref.proxy(self.__memory_manager),
                                     weakref.proxy(self.__capturer)))
        self.__define_internal_handler(self.__io.get_role(), role_type)

        return self.__io.get_role()

    @property
    def available_roles(self) -> list:
        """
        Returns available list of roles.

        Returns:
            object of list type.
        """
        return copy.deepcopy(self.__role_list)

    @property
    def opf_handler(self):
        """
        Returns current OPF handler.

        Returns:
            object of handler type.
        """
        return self.__opf_handler.handler

    @opf_handler.setter
    def opf_handler(self, value):
        """
        Set new OPF handler.

        Args:
            value - new OPF handler
        """
        self.__opf_handler.handler = value

    @property
    def _terminal(self) -> Terminal:
        return self.__terminal