Astrape/gibil/classes/sigen/modbus.py

from __future__ import annotations

from dataclasses import dataclass
from inspect import signature
import sys
from typing import Literal

try:
    from pymodbus.client import ModbusTcpClient
    from pymodbus.exceptions import ModbusException
except ImportError:  # pragma: no cover - exercised only before dependency install
    ModbusTcpClient = None  # type: ignore[assignment]

    class ModbusException(Exception):
        pass


RegisterKind = Literal["holding", "input", "coil", "discrete"]


@dataclass(frozen=True)
class ModbusReadResult:
    kind: RegisterKind
    address: int
    count: int
    values: list[int] | list[bool]


@dataclass(frozen=True)
class ModbusReadError:
    kind: RegisterKind
    address: int
    count: int
    error: str


class SigenModbusClient:
    """Small Modbus TCP client for exploring a Sigenergy plant or inverter."""

    def __init__(
        self,
        host: str,
        port: int = 502,
        unit_id: int = 1,
        timeout: float = 5.0,
        retries: int = 3,
        trace: bool = False,
    ) -> None:
        if ModbusTcpClient is None:
            raise RuntimeError(
                "pymodbus is not installed. Install dependencies with "
                "`python3 -m pip install -r requirements.txt`."
            )

        self.host = host
        self.port = port
        self.unit_id = unit_id
        self.timeout = timeout
        self.retries = retries
        self.trace = trace
        self._client = ModbusTcpClient(
            host=host,
            port=port,
            timeout=timeout,
            retries=retries,
            trace_packet=self._trace_packet if trace else None,
        )
        self._unit_keyword = self._detect_unit_keyword()

    def __enter__(self) -> SigenModbusClient:
        self.connect()
        return self

    def __exit__(self, *args: object) -> None:
        self.close()

    def connect(self) -> None:
        if not self._client.connect():
            raise ConnectionError(
                f"Could not connect to Modbus TCP target {self.host}:{self.port}"
            )

    def close(self) -> None:
        self._client.close()

    def read(
        self,
        kind: RegisterKind,
        address: int,
        count: int = 1,
    ) -> ModbusReadResult:
        if count < 1:
            raise ValueError("count must be at least 1")
        if address < 0:
            raise ValueError("address must be zero or greater")

        response = self._read_raw(kind, address, count)
        if response.isError():
            raise ModbusException(str(response))

        values = getattr(response, "registers", None)
        if values is None:
            values = getattr(response, "bits", [])
            values = list(values[:count])

        return ModbusReadResult(
            kind=kind,
            address=address,
            count=count,
            values=list(values),
        )

    def scan(
        self,
        kind: RegisterKind,
        start: int,
        count: int,
        chunk_size: int = 10,
    ) -> list[ModbusReadResult | ModbusReadError]:
        if count < 1:
            raise ValueError("count must be at least 1")
        if chunk_size < 1:
            raise ValueError("chunk_size must be at least 1")

        results: list[ModbusReadResult | ModbusReadError] = []
        stop = start + count
        address = start
        while address < stop:
            current_count = min(chunk_size, stop - address)
            try:
                results.append(self.read(kind, address, current_count))
            except Exception as exc:
                results.append(
                    ModbusReadError(
                        kind=kind,
                        address=address,
                        count=current_count,
                        error=str(exc),
                    )
                )
            address += current_count

        return results

    def _read_raw(self, kind: RegisterKind, address: int, count: int):
        if kind == "holding":
            return self._call_read(self._client.read_holding_registers, address, count)
        if kind == "input":
            return self._call_read(self._client.read_input_registers, address, count)
        if kind == "coil":
            return self._call_read(self._client.read_coils, address, count)
        if kind == "discrete":
            return self._call_read(self._client.read_discrete_inputs, address, count)

        raise ValueError(f"Unsupported register kind: {kind}")

    def _call_read(self, method, address: int, count: int):
        kwargs = {
            "address": address,
            "count": count,
            self._unit_keyword: self.unit_id,
        }
        try:
            return method(**kwargs)
        except TypeError as exc:
            if self._unit_keyword not in str(exc):
                raise

            kwargs.pop(self._unit_keyword)
            return method(address, self.unit_id, **kwargs)

    def _detect_unit_keyword(self) -> str:
        read_signature = signature(self._client.read_holding_registers)
        for keyword in ("device_id", "slave", "unit"):
            if keyword in read_signature.parameters:
                return keyword
        return "slave"

    def _trace_packet(self, sending: bool, packet: bytes) -> bytes:
        direction = "TX" if sending else "RX"
        print(f"{direction} {packet.hex(' ')}", file=sys.stderr)
        return packet