pqAutomationApp/UniTAP/dev/ports/modules/dpcd/dpcd.py

import os
from typing import Union, Tuple

from UniTAP.libs.lib_tsi.tsi import *
from UniTAP.libs.lib_tsi.tsi_io import PortIO


class DPCDRegion:

    """
    Class `DPCDRegion` describe a byte range of DPCD registers.
    Allows saving DPCD data to dpd file format `save_to_dpd`, hex files format `save_to_hex`,
    dsc file format `save_to_csv`.

    Arguments:
        base - start DPCD address of byte range.
        data - DPCD data of byte range.
    """

    def __init__(self, base: int, data: Union[bytearray, int]):
        
        if not isinstance(data, bytearray):
            self.data = bytearray()
            self.data.append(data)
        else:
            self.data = data
        self.base = base

    @property
    def size(self) -> int:
        """

        Returns length fo data.

        Returns:
            result of int object
        """
        return len(self.data)

    def save_to_dpd(self, path: str):
        """

        Save DPCD data to 'dpd' file format.

        Args:
            path (str) - full path to file

        """
        if not path.lower().endswith('.dpd'):
            path += '.DPD'
        with open(path, 'bw') as file:
            dpd_mark = 1
            file.write(dpd_mark.to_bytes(4, 'little'))
            file.write(self.base.to_bytes(4, 'little'))
            file.write(len(self.data).to_bytes(4, 'little'))
            file.write(self.data)

    def save_to_hex(self, path: str):
        """

        Save DPCD data to 'hex' file format.

        Args:
            path (str) - full path to file

        """
        if not path.lower().endswith('.hex'):
            path += '.HEX'
        with open(path, 'w') as file:
            file.write(f"\nRange 1. Start {self.base:#0{8}x}, Length {len(self.data):#0{8}x}\n")
            file.write(
                '\n'.join(
                    ' '.join(
                        f'{value:#0{2}x}' for value in self.data[i:i + 16])
                    for i in range(0, len(self.data), 16)
                )
            )

    def save_to_csv(self, path: str):
        """

        Save DPCD data to 'csv' file format.

        Args:
            path (str) - full path to file

        """
        if not path.lower().endswith('.csv'):
            path += '.csv'
        with open(path, 'w') as file:
            file.write("Address (A), Data A+0, Data A+1, Data A+2, Data A+3, Data A+4, Data A+5,"
                       " Data A+6, Data A+7\n")
            file.write(
                '\n'.join(
                    f'{self.base + 8 * i:#0{8}x},' + ','.join(f'{value:#0{2}x}'
                                                              for value in self.data[i:i + 8])
                    for i in range(0, len(self.data), 8)
                )
            )


class DPCDRegisters:

    """
    Class `DPCDRegisters` allows working with DPCD registers: writing `write` DPCD data to device,
    reading `read` DPCD data from device, loading `load_from_file` DPCD data from file.
    """

    def __init__(self, port_io: PortIO, base_ci: int, data_ci: int):
        self.__io = port_io
        self.__base_ci = base_ci
        self.__data_ci = data_ci

    def write(self, base: int, data: Union[bytearray, int, list]) -> int:
        """

        Write transferred DPCD data to device from base address.

        Args:
            base (int) - start (base) address.
            data (Union[bytearray, int]) - DPCD data.

        Returns:
            result of operation
        """
        if isinstance(data, int):
            data = bytearray(data.to_bytes(1, "big"))
        elif isinstance(data, list):
            data = bytearray(data)

        result = self.__io.set(self.__base_ci, base)
        if result < TSI_SUCCESS:
            return result
        result = self.__io.set(self.__data_ci, data, c_ubyte, len(data))
        return result

    def read(self, base: int, count: int) -> DPCDRegion:
        """

        Read DPCD data from base address in a certain quantity.

        Args:
            base (int) - start (base) address.
            count (int) - quantity of DPCD bytes.

        Returns:
            object of `DPCDRegion`
        """
        self.__io.set(self.__base_ci, base)

        result = self.__io.get(self.__data_ci, c_ubyte, count)

        return DPCDRegion(base, result[1])

    @staticmethod
    def load_from_file(path: str) -> Tuple[DPCDRegion, DPCDRegion]:
        """

        Read DPCD data from file.
        Supported formats:
        - DPD.

        Args:
            path (str) - full path to file.

        Returns:
            object of tuple with two `DPCDRegion` objects
        """
        if not os.path.exists(path):
            raise FileNotFoundError(f"DPCD load_from_file function can't find file: {path}")

        if '.DPD' not in os.path.splitext(path):
            raise TypeError("Wrong file extension. It should be .DPD")

        with open(path, 'rb') as file:
            file.read(4)
            base1 = int.from_bytes(file.read(4)[::-1], 'little')
            size1 = int.from_bytes(file.read(4)[::-1], 'little')
            data1 = bytearray(file.read(size1))
            region1 = DPCDRegion(base1, data1)

            base2 = int.from_bytes(file.read(4)[::-1], 'little')
            size2 = int.from_bytes(file.read(4)[::-1], 'little')
            data2 = bytearray(file.read(size2))
            region2 = DPCDRegion(base2, data2)

            return region1, region2