Use mypy for type checking

Author: ThyMYthOS

.github/workflows/test-python-package.yml

@@ -19,6 +19,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+    - name: Install jq
+      run: sudo apt-get update && sudo apt-get install -y jq
     - name: Install uv
       uses: astral-sh/setup-uv@e92bafb6253dcd438e0484186d7669ea7a8ca1cc # v6.4.3
       with:
@@ -27,6 +29,8 @@ jobs:
       run: uv sync --locked --all-extras --dev
     - name: Run Ruff
       run: uvx ruff check . --output-format=github
+    - name: Run type checks
+      run: uvx mypy pystiebeleltron --output=json | jq -r '"::error title=Mypy issue,file=\(.file),line=\(.line),col=\(.column)::\(.message)"'
     - name: Run tests
       run: |
         uv run pytest \

pyproject.toml

@@ -32,6 +32,7 @@ dev = [
     "pytest-asyncio>=0.26.0",
     "pytest-mock>=3.14.0",
     "ruff>=0.12.0",
+    "mypy>=1.9.0",
 ]
 
 [tool.hatch.build.targets.sdist]
@@ -61,3 +62,7 @@ testpaths = [
     "test"
 ]
 asyncio_default_fixture_loop_scope = "function"
+
+[tool.mypy]
+python_version = "3.10"
+strict = true

pystiebeleltron/__init__.py

@@ -1,14 +1,16 @@
+from __future__ import annotations
+
 import asyncio
 import logging
 from dataclasses import dataclass
 from enum import Enum
 
 from pymodbus.client import AsyncModbusTcpClient
+from pymodbus.pdu import ModbusPDU
 
 __version__ = "0.2.3"
 
-_LOGGER: logging.Logger = logging.getLogger(__package__)
-
+_LOGGER = logging.getLogger(__package__)
 
 ENERGY_DATA_BLOCK_NAME = "Energy Data"
 VIRTUAL_REGISTER_OFFSET = 100000
@@ -72,11 +74,11 @@ class ModbusRegisterBlock:
     base_address: int
     count: int
     name: str
-    registers: dict
+    registers: dict[IsgRegisters, ModbusRegister]
     register_type: RegisterType
 
 
-ENERGY_MANAGEMENT_SETTINGS_REGISTERS = {
+ENERGY_MANAGEMENT_SETTINGS_REGISTERS: dict[IsgRegisters, ModbusRegister] = {
     EnergyManagementSettingsRegisters.SWITCH_SG_READY_ON_AND_OFF: ModbusRegister(
         address=4001, name="SWITCH SG READY ON AND OFF", unit="", min=0.0, max=1.0, data_type=6, key=EnergyManagementSettingsRegisters.SWITCH_SG_READY_ON_AND_OFF
     ),
@@ -88,7 +90,7 @@ class ModbusRegisterBlock:
     ),
 }
 
-ENERGY_SYSTEM_INFORMATION_REGISTERS = {
+ENERGY_SYSTEM_INFORMATION_REGISTERS: dict[IsgRegisters, ModbusRegister] = {
     EnergySystemInformationRegisters.SG_READY_OPERATING_STATE: ModbusRegister(
         address=5001, name="SG READY OPERATING STATE", unit="", min=1.0, max=4.0, data_type=6, key=EnergySystemInformationRegisters.SG_READY_OPERATING_STATE
     ),
@@ -125,7 +127,7 @@ class ControllerModel(Enum):
     WPMsystem = 449
 
 
-async def get_controller_model(host, port) -> ControllerModel:
+async def get_controller_model(host: str, port: int) -> ControllerModel:
     """Read the model of the controller.
 
     LWA and LWZ controllers have model ids 103 and 104.
@@ -166,9 +168,9 @@ def __init__(
         self._client: AsyncModbusTcpClient = AsyncModbusTcpClient(host=host, port=port)
         self._lock = asyncio.Lock()
         self._register_blocks = register_blocks
-        self._data = {}
-        self._previous_data = {}
-        self._modbus_data = {}  # store raw data from modbus for debug purpose
+        self._data: dict[IsgRegisters, float | int | None] = {}
+        self._previous_data: dict[IsgRegisters, float | int | None] = {}
+        self._modbus_data: dict[str, ModbusPDU | None] = {}  # store raw data from modbus for debug purpose
 
     async def close(self) -> None:
         """Disconnect client."""
@@ -205,7 +207,7 @@ def get_register_descriptor(self, register: IsgRegisters) -> ModbusRegister | No
                 return descriptor
         return None
 
-    def get_register_value(self, register: IsgRegisters) -> float:
+    def get_register_value(self, register: IsgRegisters) -> float | int | None:
         """Get a value form the registers. The async_udpate needs to be called first."""
         return self._data[register]
 
@@ -222,19 +224,19 @@ async def write_register_value(self, register: IsgRegisters, value: int | float)
         else:
             raise ValueError("invalid register")
 
-    async def read_input_registers(self, device_id, address, count):
+    async def read_input_registers(self, device_id: int, address: int, count: int) -> ModbusPDU:
         """Read input registers."""
         _LOGGER.debug(f"Reading {count} input registers from {address} with device_id {device_id}")
         async with self._lock:
             return await self._client.read_input_registers(address, count=count, device_id=device_id)
 
-    async def read_holding_registers(self, device_id, address, count):
+    async def read_holding_registers(self, device_id: int, address: int, count: int) -> ModbusPDU:
         """Read holding registers."""
         _LOGGER.debug(f"Reading {count} holding registers from {address} with device_id {device_id}")
         async with self._lock:
             return await self._client.read_holding_registers(address, count=count, device_id=device_id)
 
-    def convert_value_from_modbus(self, register, register_description: ModbusRegister) -> float | int | None:
+    def convert_value_from_modbus(self, register: int, register_description: ModbusRegister) -> float | int | None:
         """Convert a modbus value to a python value."""
         if register_description.data_type == 2:
             value = self._client.convert_from_registers([register], self._client.DATATYPE.INT16)
@@ -279,9 +281,9 @@ def convert_value_to_modbus(self, value: int | float, register_description: Modb
         else:
             raise ValueError("invalid register type")
 
-    async def async_update(self):
+    async def async_update(self) -> None:
         """Request current values from heat pump."""
-        result: dict = {}
+        result: dict[IsgRegisters, float | int | None] = {}
         for registerblock in self._register_blocks:
             heat_pump_data = None
             if registerblock.register_type == RegisterType.INPUT_REGISTER:

pystiebeleltron/lwz.py

@@ -1,5 +1,7 @@
 """Modbus api for stiebel eltron heat pumps. This file is generated. Do not modify it manually."""
 
+from __future__ import annotations
+
 from enum import Enum
 
 from . import (
@@ -150,7 +152,7 @@ class LwzEnergyDataRegisters(IsgRegisters):
     ELEC_BOOSTER_DHW = 3032
 
 
-LWZ_SYSTEM_VALUES_REGISTERS = {
+LWZ_SYSTEM_VALUES_REGISTERS: dict[IsgRegisters, ModbusRegister] = {
     LwzSystemValuesRegisters.ACTUAL_ROOM_T_HC1: ModbusRegister(address=1, name="ACTUAL ROOM T HC1", unit="°C", min=-20.0, max=60.0, data_type=2, key=LwzSystemValuesRegisters.ACTUAL_ROOM_T_HC1),
     LwzSystemValuesRegisters.SET_ROOM_TEMPERATURE_HC1: ModbusRegister(
         address=2, name="SET ROOM TEMPERATURE HC1", unit="°C", min=-20.0, max=60.0, data_type=2, key=LwzSystemValuesRegisters.SET_ROOM_TEMPERATURE_HC1
@@ -211,7 +213,7 @@ class LwzEnergyDataRegisters(IsgRegisters):
     ),
 }
 
-LWZ_SYSTEM_PARAMETERS_REGISTERS = {
+LWZ_SYSTEM_PARAMETERS_REGISTERS: dict[IsgRegisters, ModbusRegister] = {
     LwzSystemParametersRegisters.OPERATING_MODE: ModbusRegister(address=1001, name="OPERATING MODE", unit="", min=0.0, max=14.0, data_type=8, key=LwzSystemParametersRegisters.OPERATING_MODE),
     LwzSystemParametersRegisters.ROOM_TEMPERATURE_DAY_HK1: ModbusRegister(
         address=1002, name="ROOM TEMPERATURE DAY", unit="°C", min=10.0, max=30.0, data_type=2, key=LwzSystemParametersRegisters.ROOM_TEMPERATURE_DAY_HK1
@@ -257,15 +259,15 @@ class LwzEnergyDataRegisters(IsgRegisters):
     LwzSystemParametersRegisters.RESTART_ISG: ModbusRegister(address=1027, name="RESTART ISG", unit="", min=0.0, max=2.0, data_type=6, key=LwzSystemParametersRegisters.RESTART_ISG),
 }
 
-LWZ_SYSTEM_STATE_REGISTERS = {
+LWZ_SYSTEM_STATE_REGISTERS: dict[IsgRegisters, ModbusRegister] = {
     LwzSystemStateRegisters.OPERATING_STATUS: ModbusRegister(address=2001, name="OPERATING STATUS", unit="", min=0.0, max=65535.0, data_type=6, key=LwzSystemStateRegisters.OPERATING_STATUS),
     LwzSystemStateRegisters.FAULT_STATUS: ModbusRegister(address=2002, name="FAULT STATUS", unit="", min=0.0, max=1.0, data_type=6, key=LwzSystemStateRegisters.FAULT_STATUS),
     LwzSystemStateRegisters.BUS_STATUS: ModbusRegister(address=2003, name="BUS STATUS", unit="", min=-4.0, max=0.0, data_type=6, key=LwzSystemStateRegisters.BUS_STATUS),
     LwzSystemStateRegisters.DEFROST_INITIATED: ModbusRegister(address=2004, name="DEFROST INITIATED", unit="", min=0.0, max=1.0, data_type=6, key=LwzSystemStateRegisters.DEFROST_INITIATED),
     LwzSystemStateRegisters.OPERATING_STATUS_2: ModbusRegister(address=2005, name="OPERATING STATUS 2", unit="", min=0.0, max=65535.0, data_type=6, key=LwzSystemStateRegisters.OPERATING_STATUS_2),
 }
 
-LWZ_ENERGY_DATA_REGISTERS = {
+LWZ_ENERGY_DATA_REGISTERS: dict[IsgRegisters, ModbusRegister] = {
     LwzEnergyDataRegisters.HEAT_METER_HTG_DAY: ModbusRegister(address=3001, name="HEAT METER HTG DAY", unit="kWh", min=0.0, max=65535.0, data_type=6, key=LwzEnergyDataRegisters.HEAT_METER_HTG_DAY),
     LwzEnergyDataRegisters.HEAT_METER_HTG_TTL_LOW: ModbusRegister(
         address=3002, name="HEAT METER HTG TTL", unit="kWh", min=0.0, max=999.0, data_type=6, key=LwzEnergyDataRegisters.HEAT_METER_HTG_TTL_LOW
@@ -364,7 +366,7 @@ def __init__(self, host: str, port: int = 502, device_id: int = 1) -> None:
             device_id,
         )
 
-    async def async_update(self):
+    async def async_update(self) -> None:
         """Request current values from heat pump."""
         await super().async_update()
         for registerblock in self._register_blocks:
@@ -402,42 +404,53 @@ async def async_update(self):
             else:
                 self._data[LwzSystemValuesRegisters.COMPRESSOR_STARTS] = compressor_starts_high * 1000 + compressor_starts_low
 
-    def get_current_temp(self):
+    def get_current_temp(self) -> float | None:
         """Get the current room temperature."""
         return self.get_register_value(LwzSystemValuesRegisters.ACTUAL_ROOM_T_HC1)
 
-    def get_target_temp(self):
+    def get_target_temp(self) -> float | None:
         """Get the target room temperature."""
         return self.get_register_value(LwzSystemParametersRegisters.ROOM_TEMPERATURE_DAY_HK1)
 
-    async def set_target_temp(self, temp):
+    async def set_target_temp(self, temp: float) -> None:
         """Set the target room temperature (day)(HC1)."""
         await self.write_register_value(LwzSystemParametersRegisters.ROOM_TEMPERATURE_DAY_HK1, temp)
 
-    def get_current_humidity(self):
+    def get_current_humidity(self) -> float | None:
         """Get the current room humidity."""
         return self.get_register_value(LwzSystemValuesRegisters.RELATIVE_HUMIDITY_HC1)
 
     # Handle operation mode
 
     def get_operation(self) -> OperatingMode:
         """Return the current mode of operation."""
-        op_mode = int(self.get_register_value(LwzSystemParametersRegisters.OPERATING_MODE))
-        return OperatingMode(op_mode)
+        op_mode = self.get_register_value(LwzSystemParametersRegisters.OPERATING_MODE)
+        if op_mode is None:
+            return OperatingMode.EMERGENCY_OPERATION
+        return OperatingMode(int(op_mode))
 
-    async def set_operation(self, mode: OperatingMode):
+    async def set_operation(self, mode: OperatingMode) -> None:
         """Set the operation mode."""
         await self.write_register_value(LwzSystemParametersRegisters.OPERATING_MODE, mode.value)
 
     def get_heating_status(self) -> bool:
         """Return heater status."""
-        return bool(int(self.get_register_value(LwzSystemStateRegisters.OPERATING_STATUS)) & (1 << 2))
+        value = self.get_register_value(LwzSystemStateRegisters.OPERATING_STATUS)
+        if value is None:
+            return False
+        return bool(int(value) & (1 << 2))
 
     def get_cooling_status(self) -> bool:
         """Cooling status."""
-        return bool(int(self.get_register_value(LwzSystemStateRegisters.OPERATING_STATUS)) & (1 << 3))
+        value = self.get_register_value(LwzSystemStateRegisters.OPERATING_STATUS)
+        if value is None:
+            return False
+        return bool(int(value) & (1 << 3))
 
     def get_filter_alarm_status(self) -> bool:
         """Return filter alarm."""
+        value = self.get_register_value(LwzSystemStateRegisters.OPERATING_STATUS)
+        if value is None:
+            return False
         filter_mask = (1 << 8) | (1 << 12) | (1 << 13)
-        return bool(int(self.get_register_value(LwzSystemStateRegisters.OPERATING_STATUS)) & filter_mask)
+        return bool(int(value) & filter_mask)

pystiebeleltron/pystiebeleltron.py

@@ -17,7 +17,10 @@
 8    | 0 to 255   | 1           | 1           | No     | 1      | 5
 """
 
+from __future__ import annotations
+
 from pymodbus.client.mixin import ModbusClientMixin
+from pymodbus.pdu import ModbusPDU
 
 # Error - sensor lead is missing or disconnected.
 ERROR_NOTAVAILABLE = -60
@@ -156,7 +159,7 @@
 class StiebelEltronAPI:
     """Stiebel Eltron API."""
 
-    def __init__(self, conn: ModbusClientMixin, device_id=1, update_on_read=False):
+    def __init__(self, conn: ModbusClientMixin[ModbusPDU], device_id: int = 1, update_on_read: bool = False) -> None:
         """Initialize Stiebel Eltron communication."""
         self._conn = conn
         self._block_1_input_regs = B1_REGMAP_INPUT
@@ -165,7 +168,7 @@ def __init__(self, conn: ModbusClientMixin, device_id=1, update_on_read=False):
         self._device_id = device_id
         self._update_on_read = update_on_read
 
-    def update(self):
+    def update(self) -> bool:
         """Request current values from heat pump."""
         ret = True
         try:
@@ -189,7 +192,7 @@ def update(self):
 
         return ret
 
-    def get_conv_val(self, name: str):
+    def get_conv_val(self, name: str) -> float | int | None:
         """Read and convert value.
 
         Args:
@@ -234,60 +237,72 @@ def get_conv_val(self, name: str):
 
     # Handle room temperature & humidity
 
-    def get_current_temp(self):
+    def get_current_temp(self) -> float | None:
         """Get the current room temperature."""
         if self._update_on_read:
             self.update()
         return self.get_conv_val("ACTUAL_ROOM_TEMPERATURE_HC1")
 
-    def get_target_temp(self):
+    def get_target_temp(self) -> float | None:
         """Get the target room temperature."""
         if self._update_on_read:
             self.update()
         return self.get_conv_val("ROOM_TEMP_HEAT_DAY_HC1")
 
-    def set_target_temp(self, temp: float):
+    def set_target_temp(self, temp: float) -> None:
         """Set the target room temperature (day)(HC1)."""
         self._conn.write_register(device_id=self._device_id, address=(self._block_2_holding_regs["ROOM_TEMP_HEAT_DAY_HC1"]["addr"]), value=round(temp * 10.0))
 
-    def get_current_humidity(self):
+    def get_current_humidity(self) -> float | None:
         """Get the current room humidity."""
         if self._update_on_read:
             self.update()
         return self.get_conv_val("RELATIVE_HUMIDITY_HC1")
 
     # Handle operation mode
 
-    def get_operation(self):
+    def get_operation(self) -> str:
         """Return the current mode of operation."""
         if self._update_on_read:
             self.update()
 
         op_mode = self.get_conv_val("OPERATING_MODE")
-        return B2_OPERATING_MODE_READ.get(op_mode, "UNKNOWN")
+        return B2_OPERATING_MODE_READ.get(int(op_mode) if op_mode is not None else -1, "UNKNOWN")
 
-    def set_operation(self, mode: str):
+    def set_operation(self, mode: str) -> None:
         """Set the operation mode."""
-        self._conn.write_register(device_id=self._device_id, address=(self._block_2_holding_regs["OPERATING_MODE"]["addr"]), value=B2_OPERATING_MODE_WRITE.get(mode))
+        value = B2_OPERATING_MODE_WRITE.get(mode)
+        if value is None:
+            raise ValueError(f"Invalid operation mode: {mode}")
+        self._conn.write_register(device_id=self._device_id, address=(self._block_2_holding_regs["OPERATING_MODE"]["addr"]), value=value)
 
     # Handle device status
 
-    def get_heating_status(self):
+    def get_heating_status(self) -> bool:
         """Return heater status."""
         if self._update_on_read:
             self.update()
-        return bool(self.get_conv_val("OPERATING_STATUS") & B3_OPERATING_STATUS["HEATING"])
+        val = self.get_conv_val("OPERATING_STATUS")
+        if val is None:
+            return False
+        return bool(int(val) & B3_OPERATING_STATUS["HEATING"])
 
-    def get_cooling_status(self):
+    def get_cooling_status(self) -> bool:
         """Cooling status."""
         if self._update_on_read:
             self.update()
-        return bool(self.get_conv_val("OPERATING_STATUS") & B3_OPERATING_STATUS["COOLING"])
+        val = self.get_conv_val("OPERATING_STATUS")
+        if val is None:
+            return False
+        return bool(int(val) & B3_OPERATING_STATUS["COOLING"])
 
-    def get_filter_alarm_status(self):
+    def get_filter_alarm_status(self) -> bool:
         """Return filter alarm."""
         if self._update_on_read:
             self.update()
 
         filter_mask = B3_OPERATING_STATUS["FILTER"] | B3_OPERATING_STATUS["FILTER_EXTRACT_AIR"] | B3_OPERATING_STATUS["FILTER_VENTILATION_AIR"]
-        return bool(self.get_conv_val("OPERATING_STATUS") & filter_mask)
+        val = self.get_conv_val("OPERATING_STATUS")
+        if val is None:
+            return False
+        return bool(int(val) & filter_mask)