From 11409115c007a85c17d30c678baed922a2d7f634 Mon Sep 17 00:00:00 2001
From: mvn23 <schopdiedwaas@gmail.com>
Date: Tue, 11 Feb 2025 13:04:13 +0100
Subject: [PATCH] Fix type annotations and linter errors

---
 custom_components/saj_modbus/hub.py    | 34 +++++++++++++++-----------
 custom_components/saj_modbus/number.py |  7 ++++--
 2 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/custom_components/saj_modbus/hub.py b/custom_components/saj_modbus/hub.py
index af2942f..2ad3472 100644
--- a/custom_components/saj_modbus/hub.py
+++ b/custom_components/saj_modbus/hub.py
@@ -13,7 +13,8 @@
 from homeassistant.components.number import DOMAIN as NUMBER_DOMAIN
 from pymodbus.client import ModbusTcpClient
 from pymodbus.constants import Endian
-from pymodbus.exceptions import ConnectionException
+from pymodbus.exceptions import ConnectionException, ModbusException
+from pymodbus.pdu import ModbusPDU
 from .payload import BinaryPayloadDecoder
 
 from .const import (
@@ -63,7 +64,7 @@ def close(self) -> None:
         """Disconnect client."""
         with self._lock:
             self._client.close()
-            
+
     def _read_holding_registers(self, unit, address, count):
         """Read holding registers."""
         with self._lock:
@@ -71,13 +72,13 @@ def _read_holding_registers(self, unit, address, count):
                 address=address, count=count, slave=unit
             )
 
-    def _write_registers(self, unit: int, address: int, values: list[int] | int):
+    def _write_registers(self, unit: int, address: int, values: list[int]) -> ModbusPDU:
         """Write registers."""
         with self._lock:
             return self._client.write_registers(
                 address=address, values=values, slave=unit
             )
-            
+
     async def _async_update_data(self) -> dict:
         realtime_data = {}
         try:
@@ -304,25 +305,30 @@ def translate_fault_code_to_messages(
 
         return messages
 
-    def limiter_is_disabled(self):
+    def limiter_is_disabled(self) -> bool:
         """Return True if the limiter entity is disabled, False otherwise."""
         ent_reg = entity_registry.async_get(self.hass)
-        limiter_entity_id = ent_reg.async_get_entity_id(NUMBER_DOMAIN, DOMAIN, f"{self.name}_limitpower")
-        if limiter_entity_id is None:
+        limiter_entity_id = ent_reg.async_get_entity_id(
+            NUMBER_DOMAIN, DOMAIN, f"{self.name}_limitpower"
+        )
+        if (
+            limiter_entity_id is None
+            or (ent_reg_entry := ent_reg.async_get(limiter_entity_id)) is None
+        ):
             return True
-        return ent_reg.async_get(limiter_entity_id).disabled
+        return ent_reg_entry.disabled
 
-    def set_limitpower(self, value: int):
+    def set_limitpower(self, value: float) -> None:
         """Limit the power output of the inverter."""
         if self.limiter_is_disabled():
             return
-        response = self._write_registers(unit=1, address=0x801F, values=[int(value*10)])
+        response = self._write_registers(unit=1, address=0x801F, values=[int(value * 10)])
         if response.isError():
             return
         self.data["limitpower"] = value
         self.hass.add_job(self.async_update_listeners)
 
-    def set_date_and_time(self, date_time: datetime = None):
+    def set_date_and_time(self, date_time: datetime | None = None) -> None:
         """Set the time and date on the inverter."""
         if date_time is None:
             date_time = datetime.now()
@@ -331,14 +337,14 @@ def set_date_and_time(self, date_time: datetime = None):
             date_time.year,
             (date_time.month << 8) + date_time.day,
             (date_time.hour << 8) + date_time.minute,
-            (date_time.second << 8)
+            (date_time.second << 8),
         ]
 
         response = self._write_registers(unit=1, address=0x8020, values=values)
         if response.isError():
-            raise response
+            raise ModbusException("Error setting date and time")
 
-    def set_value(self, key: str, value: int):
+    def set_value(self, key: str, value: float) -> None:
         """Set value matching key."""
         if key == "limitpower":
             self.set_limitpower(value)
diff --git a/custom_components/saj_modbus/number.py b/custom_components/saj_modbus/number.py
index 1593083..d3b22b7 100644
--- a/custom_components/saj_modbus/number.py
+++ b/custom_components/saj_modbus/number.py
@@ -1,4 +1,5 @@
 """Number entity for SAJ Modbus integration."""
+
 from __future__ import annotations
 
 import logging
@@ -46,6 +47,8 @@ async def async_setup_entry(hass, entry, async_add_entities):
 class SajNumber(CoordinatorEntity, NumberEntity):
     """Representation of an SAJ Modbus number."""
 
+    coordinator: SAJModbusHub
+
     def __init__(
         self,
         platform_name: str,
@@ -63,14 +66,14 @@ def __init__(
 
     @property
     def available(self) -> bool:
+        """Return entity availability."""
         return self.native_value is not None
 
     @property
-    def native_value(self) -> Optional[int]:
+    def native_value(self) -> float | None:
         """Return the state of the number entity."""
         return self.coordinator.data.get(self.entity_description.key, None)
 
     def set_native_value(self, value: float) -> None:
         """Update the current value."""
         self.coordinator.set_value(self.entity_description.key, value)
-