Improve type hint in compensation sensor entity (#77027)

homeassistant/components/compensation/sensor.py

@@ -2,6 +2,9 @@
 from __future__ import annotations
 
 import logging
+from typing import Any
+
+import numpy as np
 
 from homeassistant.components.sensor import SensorEntity
 from homeassistant.const import (
@@ -12,7 +15,7 @@ from homeassistant.const import (
     CONF_UNIT_OF_MEASUREMENT,
     STATE_UNKNOWN,
 )
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import Event, HomeAssistant, State, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_state_change_event
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
@@ -42,11 +45,11 @@ async def async_setup_platform(
     if discovery_info is None:
         return
 
-    compensation = discovery_info[CONF_COMPENSATION]
+    compensation: str = discovery_info[CONF_COMPENSATION]
-    conf = hass.data[DATA_COMPENSATION][compensation]
+    conf: dict[str, Any] = hass.data[DATA_COMPENSATION][compensation]
 
-    source = conf[CONF_SOURCE]
+    source: str = conf[CONF_SOURCE]
-    attribute = conf.get(CONF_ATTRIBUTE)
+    attribute: str | None = conf.get(CONF_ATTRIBUTE)
     name = f"{DEFAULT_NAME} {source}"
     if attribute is not None:
         name = f"{name} {attribute}"
@@ -69,26 +72,27 @@ async def async_setup_platform(
 class CompensationSensor(SensorEntity):
     """Representation of a Compensation sensor."""
 
+    _attr_should_poll = False
+
     def __init__(
         self,
-        unique_id,
+        unique_id: str | None,
-        name,
+        name: str,
-        source,
+        source: str,
-        attribute,
+        attribute: str | None,
-        precision,
+        precision: int,
-        polynomial,
+        polynomial: np.poly1d,
-        unit_of_measurement,
+        unit_of_measurement: str | None,
-    ):
+    ) -> None:
         """Initialize the Compensation sensor."""
         self._source_entity_id = source
         self._precision = precision
         self._source_attribute = attribute
-        self._unit_of_measurement = unit_of_measurement
+        self._attr_native_unit_of_measurement = unit_of_measurement
         self._poly = polynomial
         self._coefficients = polynomial.coefficients.tolist()
-        self._state = None
+        self._attr_unique_id = unique_id
-        self._unique_id = unique_id
+        self._attr_name = name
-        self._name = name
 
     async def async_added_to_hass(self) -> None:
         """Handle added to Hass."""
@@ -101,27 +105,7 @@ class CompensationSensor(SensorEntity):
         )
 
     @property
-    def unique_id(self):
+    def extra_state_attributes(self) -> dict[str, Any]:
-        """Return the unique id of this sensor."""
-        return self._unique_id
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return self._name
-
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._state
-
-    @property
-    def extra_state_attributes(self):
         """Return the state attributes of the sensor."""
         ret = {
             ATTR_SOURCE: self._source_entity_id,
@@ -131,33 +115,27 @@ class CompensationSensor(SensorEntity):
             ret[ATTR_SOURCE_ATTRIBUTE] = self._source_attribute
         return ret
 
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit the value is expressed in."""
-        return self._unit_of_measurement
-
     @callback
-    def _async_compensation_sensor_state_listener(self, event):
+    def _async_compensation_sensor_state_listener(self, event: Event) -> None:
         """Handle sensor state changes."""
+        new_state: State | None
         if (new_state := event.data.get("new_state")) is None:
             return
 
-        if self._unit_of_measurement is None and self._source_attribute is None:
+        if self.native_unit_of_measurement is None and self._source_attribute is None:
-            self._unit_of_measurement = new_state.attributes.get(
+            self._attr_native_unit_of_measurement = new_state.attributes.get(
                 ATTR_UNIT_OF_MEASUREMENT
             )
 
+        if self._source_attribute:
+            value = new_state.attributes.get(self._source_attribute)
+        else:
+            value = None if new_state.state == STATE_UNKNOWN else new_state.state
         try:
-            if self._source_attribute:
+            self._attr_native_value = round(self._poly(float(value)), self._precision)
-                value = float(new_state.attributes.get(self._source_attribute))
-            else:
-                value = (
-                    None if new_state.state == STATE_UNKNOWN else float(new_state.state)
-                )
-            self._state = round(self._poly(value), self._precision)
 
         except (ValueError, TypeError):
-            self._state = None
+            self._attr_native_value = None
             if self._source_attribute:
                 _LOGGER.warning(
                     "%s attribute %s is not numerical",