Merge pull request #109 from klejejs/chore/type-improvements

Add type improvements, lint fixes

custom_components/thermia/binary_sensor.py

@@ -2,20 +2,21 @@
 
 from __future__ import annotations
 from typing import List
+from ThermiaOnlineAPI import ThermiaHeatPump
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.components.binary_sensor import BinarySensorDeviceClass
+
 from .binary_sensors.operational_or_power_status_binary_sensor import (
     ThermiaOperationalOrPowerStatusBinarySensor,
 )
-from ThermiaOnlineAPI import ThermiaHeatPump
-
 from .const import (
     DOMAIN,
     MDI_INFORMATION_OUTLINE_ICON,
 )
+from .coordinator import ThermiaDataUpdateCoordinator
 
 
 async def async_setup_entry(
@@ -25,14 +26,17 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Thermia binary sensors."""
 
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: ThermiaDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     hass_thermia_binary_sensors = []
 
     heat_pumps: List[ThermiaHeatPump] = coordinator.data.heat_pumps
 
     for idx, heat_pump in enumerate(heat_pumps):
-        if heat_pump.available_operational_statuses is not None and heat_pump.running_operational_statuses is not None:
+        if (
+            heat_pump.available_operational_statuses is not None
+            and heat_pump.running_operational_statuses is not None
+        ):
             for operational_status in heat_pump.available_operational_statuses:
                 name = operational_status.replace("_", " ").title()
                 hass_thermia_binary_sensors.append(
@@ -44,11 +48,14 @@ async def async_setup_entry(
                         MDI_INFORMATION_OUTLINE_ICON,
                         BinarySensorDeviceClass.RUNNING,
                         operational_status,
-                        "running_operational_statuses"
+                        "running_operational_statuses",
                     )
                 )
 
-        if heat_pump.available_power_statuses is not None and heat_pump.running_power_statuses is not None:
+        if (
+            heat_pump.available_power_statuses is not None
+            and heat_pump.running_power_statuses is not None
+        ):
             for power_status in heat_pump.available_power_statuses:
                 name = power_status.replace("_", " ").title()
                 hass_thermia_binary_sensors.append(
@@ -60,7 +67,7 @@ async def async_setup_entry(
                         MDI_INFORMATION_OUTLINE_ICON,
                         BinarySensorDeviceClass.RUNNING,
                         power_status,
-                        "running_power_statuses"
+                        "running_power_statuses",
                     )
                 )
 

custom_components/thermia/binary_sensors/operational_or_power_status_binary_sensor.py

@@ -2,35 +2,41 @@
 
 from __future__ import annotations
 
-from homeassistant.components.binary_sensor import BinarySensorEntity
+from homeassistant.components.binary_sensor import (
+    BinarySensorDeviceClass,
+    BinarySensorEntity,
+)
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from ..const import DOMAIN
+from ..coordinator import ThermiaDataUpdateCoordinator
 
 
-class ThermiaOperationalOrPowerStatusBinarySensor(CoordinatorEntity, BinarySensorEntity):
+class ThermiaOperationalOrPowerStatusBinarySensor(
+    CoordinatorEntity[ThermiaDataUpdateCoordinator], BinarySensorEntity
+):
     """Representation of an Thermia Operational or Power Status binary sensor."""
 
     def __init__(
         self,
         coordinator,
-        idx,
-        is_online_prop,
-        binary_sensor_name,
-        mdi_icon,
-        device_class,
-        status_value,
-        running_status_list,
+        idx: int,
+        is_online_prop: str,
+        binary_sensor_name: str,
+        mdi_icon: str,
+        device_class: BinarySensorDeviceClass,
+        status_value: str,
+        running_status_list: str,
     ):
         super().__init__(coordinator)
-        self.idx = idx
+        self.idx: int = idx
 
-        self._is_online_prop = is_online_prop
-        self._binary_sensor_name = binary_sensor_name
-        self._mdi_icon = mdi_icon
-        self._device_class = device_class
-        self._status_value = status_value
-        self._running_status_list = running_status_list
+        self._is_online_prop: str = is_online_prop
+        self._binary_sensor_name: str = binary_sensor_name
+        self._mdi_icon: str = mdi_icon
+        self._device_class: BinarySensorDeviceClass = device_class
+        self._status_value: str = status_value
+        self._running_status_list: str = running_status_list
 
     @property
     def available(self):

custom_components/thermia/sensor.py

@@ -17,6 +17,7 @@
     MDI_TEMPERATURE_ICON,
     MDI_TIMER_COG_OUTLINE_ICON,
 )
+from .coordinator import ThermiaDataUpdateCoordinator
 
 
 async def async_setup_entry(
@@ -393,7 +394,4 @@ async def async_setup_entry(
         for idx, _ in enumerate(coordinator.data.heat_pumps)
     ]
 
-    async_add_entities([
-        *hass_thermia_active_alarms_sensors,
-        *hass_thermia_sensors
-    ])
+    async_add_entities([*hass_thermia_active_alarms_sensors, *hass_thermia_sensors])

custom_components/thermia/sensors/active_alarms_sensor.py

@@ -6,14 +6,17 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from ..const import DOMAIN
+from ..coordinator import ThermiaDataUpdateCoordinator
 
 
-class ThermiaActiveAlarmsSensor(CoordinatorEntity, SensorEntity):
+class ThermiaActiveAlarmsSensor(
+    CoordinatorEntity[ThermiaDataUpdateCoordinator], SensorEntity
+):
     """Representation of an Thermia active alarms sensor."""
 
-    def __init__(self, coordinator, idx):
+    def __init__(self, coordinator, idx: int):
         super().__init__(coordinator)
-        self.idx = idx
+        self.idx: int = idx
 
     @property
     def available(self):

custom_components/thermia/sensors/generic_sensor.py

@@ -6,35 +6,38 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from ..const import DOMAIN
+from ..coordinator import ThermiaDataUpdateCoordinator
 
 
-class ThermiaGenericSensor(CoordinatorEntity, SensorEntity):
+class ThermiaGenericSensor(
+    CoordinatorEntity[ThermiaDataUpdateCoordinator], SensorEntity
+):
     """Representation of an Thermia generic sensor."""
 
     def __init__(
         self,
         coordinator,
         idx,
-        is_online_prop,
-        sensor_name,
-        mdi_icon,
-        entity_category,
-        device_class,
-        state_class,
-        value_prop,
-        unit_of_measurement,
+        is_online_prop: str,
+        sensor_name: str,
+        mdi_icon: str,
+        entity_category: str,
+        device_class: str | None,
+        state_class: str,
+        value_prop: str,
+        unit_of_measurement: str | None,
     ):
         super().__init__(coordinator)
-        self.idx = idx
-
-        self._is_online_prop = is_online_prop
-        self._sensor_name = sensor_name
-        self._mdi_icon = mdi_icon
-        self._entity_category = entity_category
-        self._device_class = device_class
-        self._state_class = state_class
-        self._value_prop = value_prop
-        self._unit_of_measurement = unit_of_measurement
+        self.idx: int = idx
+
+        self._is_online_prop: str = is_online_prop
+        self._sensor_name: str = sensor_name
+        self._mdi_icon: str = mdi_icon
+        self._entity_category: str = entity_category
+        self._device_class: str | None = device_class
+        self._state_class: str = state_class
+        self._value_prop: str = value_prop
+        self._unit_of_measurement: str | None = unit_of_measurement
 
     @property
     def available(self):

custom_components/thermia/services.py

@@ -9,8 +9,8 @@
 from homeassistant.helpers.service import async_extract_entity_ids
 from homeassistant.helpers.template import device_attr
 
-from .const import DEBUG_ACTION_NAME, DEFAULT_DEBUG_FILENAME, DOMAIN
 from .coordinator import ThermiaDataUpdateCoordinator
+from .const import DEBUG_ACTION_NAME, DEFAULT_DEBUG_FILENAME, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -35,10 +35,7 @@ async def async_handle_heat_pump_debug(self, call: ServiceCall):
         entity_ids = await async_extract_entity_ids(self.hass, call)
 
         entity_ids = list(
-            filter(
-                lambda entity_id: entity_id.startswith("water_heater."),
-                entity_ids
-            )
+            filter(lambda entity_id: entity_id.startswith("water_heater."), entity_ids)
         )
 
         if len(entity_ids) == 0 or len(entity_ids) > 1:
@@ -74,9 +71,7 @@ async def async_handle_heat_pump_debug(self, call: ServiceCall):
         )
 
         if heat_pump is None:
-            raise ServiceValidationError(
-                "Cannot find heat pump by unique_id"
-            )
+            raise ServiceValidationError("Cannot find heat pump by unique_id")
 
         debug_data = await self.hass.async_add_executor_job(lambda: heat_pump.debug())
 

custom_components/thermia/switch.py

@@ -6,11 +6,11 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from .const import DOMAIN
+from .coordinator import ThermiaDataUpdateCoordinator
 from .switches.hot_water_switch import ThermiaHotWaterSwitch
 from .switches.hot_water_boost_switch import ThermiaHotWaterBoostSwitch
 
-from .const import DOMAIN
-
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -19,7 +19,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Thermia switches."""
 
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: ThermiaDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     hass_thermia_switches = []
 

custom_components/thermia/switches/hot_water_boost_switch.py

@@ -6,14 +6,17 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from ..const import DOMAIN
+from ..coordinator import ThermiaDataUpdateCoordinator
 
 
-class ThermiaHotWaterBoostSwitch(CoordinatorEntity, SwitchEntity):
+class ThermiaHotWaterBoostSwitch(
+    CoordinatorEntity[ThermiaDataUpdateCoordinator], SwitchEntity
+):
     """Representation of an Thermia hot water boost switch."""
 
-    def __init__(self, coordinator, idx):
+    def __init__(self, coordinator, idx: int):
         super().__init__(coordinator)
-        self.idx = idx
+        self.idx: int = idx
 
     @property
     def available(self):

custom_components/thermia/switches/hot_water_switch.py

@@ -6,14 +6,17 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from ..const import DOMAIN
+from ..coordinator import ThermiaDataUpdateCoordinator
 
 
-class ThermiaHotWaterSwitch(CoordinatorEntity, SwitchEntity):
+class ThermiaHotWaterSwitch(
+    CoordinatorEntity[ThermiaDataUpdateCoordinator], SwitchEntity
+):
     """Representation of an Thermia hot water switch."""
 
-    def __init__(self, coordinator, idx):
+    def __init__(self, coordinator, idx: int):
         super().__init__(coordinator)
-        self.idx = idx
+        self.idx: int = idx
 
     @property
     def available(self):

custom_components/thermia/water_heater.py

@@ -39,7 +39,9 @@ async def async_setup_entry(
     async_add_entities(hass_water_heaters)
 
 
-class ThermiaWaterHeater(CoordinatorEntity[ThermiaDataUpdateCoordinator], WaterHeaterEntity):
+class ThermiaWaterHeater(
+    CoordinatorEntity[ThermiaDataUpdateCoordinator], WaterHeaterEntity
+):
     """Representation of an Thermia water heater."""
 
     def __init__(self, coordinator: ThermiaDataUpdateCoordinator, idx: int):
@@ -135,7 +137,11 @@ def supported_features(self):
         """Return the list of supported features."""
         features = WaterHeaterEntityFeature.TARGET_TEMPERATURE
 
-        if self.current_operation is not None and self.coordinator.data.heat_pumps[self.idx].is_operation_mode_read_only is False:
+        if (
+            self.current_operation is not None
+            and self.coordinator.data.heat_pumps[self.idx].is_operation_mode_read_only
+            is False
+        ):
             features |= WaterHeaterEntityFeature.OPERATION_MODE
 
         return features