Typing

andrew-codechimp · Oct 27, 2024 · 7946837 · 7946837
1 parent 19cd9ce
commit 7946837
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 18 deletions.
diff --git a/custom_components/hive_local_thermostat/entity.py b/custom_components/hive_local_thermostat/entity.py
@@ -17,7 +17,6 @@ class HiveEntityDescription(EntityDescription):
     entity_id: str | None = None
     topic: str
     entry_id: str | None = None
-    icons_by_state: dict[str, str] | None = None
     model: str | None = None
 
 class HiveEntity(Entity):

diff --git a/custom_components/hive_local_thermostat/number.py b/custom_components/hive_local_thermostat/number.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from datetime import datetime
 
 from homeassistant.components.number import (
     NumberDeviceClass,
@@ -153,6 +154,7 @@ class HiveNumber(HiveEntity, RestoreNumber):
 
     entity_description: HiveNumberEntityDescription
     _state: float | None
+    _last_updated: datetime | None
 
     def __init__(
         self,
@@ -167,7 +169,6 @@ def __init__(
         self._attr_has_entity_name = True
         self._topic = entity_description.topic
         self._state = None
-        self._attributes = {}
         self._last_updated = None
 
         super().__init__(entity_description)
@@ -214,11 +215,6 @@ async def async_set_native_value(self, value: float) -> None:
 
         self.async_write_ha_state()
 
-    @property
-    def extra_state_attributes(self):
-        """Attributes of the sensor."""
-        return self._attributes
-
     def process_update(self, mqtt_data) -> None:
         """Update the state of the sensor."""
         if (

diff --git a/custom_components/hive_local_thermostat/select.py b/custom_components/hive_local_thermostat/select.py
@@ -80,9 +80,10 @@ def __init__(
         self._attr_unique_id = f"{DOMAIN}_{entity_description.name}_{entity_description.key}".lower()
         self._attr_has_entity_name = True
         self._topic = entity_description.topic
-        self._attr_options = entity_description.options
         self._attr_current_option = None
         self._mqtt_data = None
+        if entity_description.options:
+            self._attr_options = entity_description.options
 
         super().__init__(entity_description)
 

diff --git a/custom_components/hive_local_thermostat/sensor.py b/custom_components/hive_local_thermostat/sensor.py
@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
+from typing import Any, cast
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -36,7 +38,9 @@ class HiveSensorEntityDescription(
 ):
     """Class describing Hive sensor entities."""
 
-    func: any | None = None
+    icons_by_state: dict[str, str] | None = None
+    value_fn: Callable[[dict[str, Any]], str | int | float | None]
+    # value_fn: any | None = None
     running_state: bool = False
 
 
@@ -60,7 +64,8 @@ async def async_setup_entry(
                     "preheating": "mdi:radiator",
                 },
                 name=config_entry.title,
-                func=lambda js: js["running_state_heat"],
+                value_fn=lambda data: cast(str, data["running_state_heat"]),
+                # value_fn=lambda js: js["running_state_heat"],
                 topic=config_entry.options[CONF_MQTT_TOPIC],
                 entry_id=config_entry.entry_id,
                 model=config_entry.options[CONF_MODEL],
@@ -74,7 +79,8 @@ async def async_setup_entry(
                 device_class=SensorDeviceClass.TEMPERATURE,
                 native_unit_of_measurement=UnitOfTemperature.CELSIUS,
                 suggested_display_precision=1,
-                func=lambda js: js["local_temperature_heat"],
+                value_fn=lambda data: cast(float, data["local_temperature_heat"]),
+                # value_fn=lambda js: js["local_temperature_heat"],
                 topic=config_entry.options[CONF_MQTT_TOPIC],
                 entry_id=config_entry.entry_id,
                 model=config_entry.options[CONF_MODEL],
@@ -89,7 +95,8 @@ async def async_setup_entry(
                     "off": "mdi:water-boiler-off",
                 },
                 name=config_entry.title,
-                func=lambda js: js["running_state_water"],
+                value_fn=lambda data: cast(str, data["running_state_water"]),
+                # value_fn=lambda js: js["running_state_water"],
                 topic=config_entry.options[CONF_MQTT_TOPIC],
                 entry_id=config_entry.entry_id,
                 model=config_entry.options[CONF_MODEL],
@@ -109,7 +116,8 @@ async def async_setup_entry(
                     "preheating": "mdi:radiator",
                 },
                 name=config_entry.title,
-                func=lambda js: js["running_state"],
+                value_fn=lambda data: cast(str, data["running_state"]),
+                # value_fn=lambda js: js["running_state"],
                 topic=config_entry.options[CONF_MQTT_TOPIC],
                 entry_id=config_entry.entry_id,
                 model=config_entry.options[CONF_MODEL],
@@ -123,7 +131,8 @@ async def async_setup_entry(
                 device_class=SensorDeviceClass.TEMPERATURE,
                 native_unit_of_measurement=UnitOfTemperature.CELSIUS,
                 suggested_display_precision=1,
-                func=lambda js: js["local_temperature"],
+                value_fn=lambda data: cast(float, data["local_temperature"]),
+                # value_fn=lambda js: js["local_temperature"],
                 topic=config_entry.options[CONF_MQTT_TOPIC],
                 entry_id=config_entry.entry_id,
                 model=config_entry.options[CONF_MODEL],
@@ -158,7 +167,7 @@ def __init__(
             f"{DOMAIN}_{entity_description.name}_{entity_description.key}".lower()
         )
         self._attr_has_entity_name = True
-        self._func = entity_description.func
+        self._func = entity_description.value_fn
         self._topic = entity_description.topic
 
         super().__init__(entity_description)
@@ -174,9 +183,12 @@ def process_update(self, mqtt_data) -> None:
         if self.entity_description.running_state:
             if new_value == "":
                 new_value = "preheating"
-            self._attr_icon = self.entity_description.icons_by_state.get(
-                new_value, ICON_UNKNOWN
-            )
+            if self.entity_description.icons_by_state:
+                self._attr_icon = self.entity_description.icons_by_state.get(
+                    cast(str, new_value), ICON_UNKNOWN
+                )
+            else:
+                self._attr_icon = ICON_UNKNOWN
 
         if self.entity_description.device_class == SensorDeviceClass.TEMPERATURE:
             new_value = show_temp(