Improve type hint in derivative sensor entity (#77038)

homeassistant/components/derivative/sensor.py

@@ -1,7 +1,7 @@
 """Numeric derivative of data coming from a source sensor over time."""
 from __future__ import annotations
 
-from datetime import timedelta
+from datetime import datetime, timedelta
 from decimal import Decimal, DecimalException
 import logging
 
@@ -20,7 +20,7 @@ from homeassistant.const import (
     TIME_MINUTES,
     TIME_SECONDS,
 )
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import Event, HomeAssistant, State, callback
 from homeassistant.helpers import config_validation as cv, entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_state_change_event
@@ -133,6 +133,9 @@ async def async_setup_platform(
 class DerivativeSensor(RestoreEntity, SensorEntity):
     """Representation of an derivative sensor."""
 
+    _attr_icon = ICON
+    _attr_should_poll = False
+
     def __init__(
         self,
         *,
@@ -150,19 +153,19 @@ class DerivativeSensor(RestoreEntity, SensorEntity):
         self._sensor_source_id = source_entity
         self._round_digits = round_digits
         self._state = 0
-        self._state_list = (
+        # List of tuples with (timestamp_start, timestamp_end, derivative)
-            []
+        self._state_list: list[tuple[datetime, datetime, Decimal]] = []
-        )  # List of tuples with (timestamp_start, timestamp_end, derivative)
 
-        self._name = name if name is not None else f"{source_entity} derivative"
+        self._attr_name = name if name is not None else f"{source_entity} derivative"
+        self._attr_extra_state_attributes = {ATTR_SOURCE_ID: source_entity}
 
         if unit_of_measurement is None:
             final_unit_prefix = "" if unit_prefix is None else unit_prefix
             self._unit_template = f"{final_unit_prefix}{{}}/{unit_time}"
             # we postpone the definition of unit_of_measurement to later
-            self._unit_of_measurement = None
+            self._attr_native_unit_of_measurement = None
         else:
-            self._unit_of_measurement = unit_of_measurement
+            self._attr_native_unit_of_measurement = unit_of_measurement
 
         self._unit_prefix = UNIT_PREFIXES[unit_prefix]
         self._unit_time = UNIT_TIME[unit_time]
@@ -178,20 +181,21 @@ class DerivativeSensor(RestoreEntity, SensorEntity):
                 _LOGGER.warning("Could not restore last state: %s", err)
 
         @callback
-        def calc_derivative(event):
+        def calc_derivative(event: Event) -> None:
             """Handle the sensor state changes."""
-            old_state = event.data.get("old_state")
+            old_state: State | None
-            new_state = event.data.get("new_state")
+            new_state: State | None
             if (
-                old_state is None
+                (old_state := event.data.get("old_state")) is None
                 or old_state.state in (STATE_UNKNOWN, STATE_UNAVAILABLE)
+                or (new_state := event.data.get("new_state")) is None
                 or new_state.state in (STATE_UNKNOWN, STATE_UNAVAILABLE)
             ):
                 return
 
-            if self._unit_of_measurement is None:
+            if self.native_unit_of_measurement is None:
                 unit = new_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
-                self._unit_of_measurement = self._unit_template.format(
+                self._attr_native_unit_of_measurement = self._unit_template.format(
                     "" if unit is None else unit
                 )
 
@@ -242,7 +246,7 @@ class DerivativeSensor(RestoreEntity, SensorEntity):
             if elapsed_time > self._time_window:
                 derivative = new_derivative
             else:
-                derivative = 0
+                derivative = Decimal(0)
                 for (start, end, value) in self._state_list:
                     weight = calculate_weight(start, end, new_state.last_updated)
                     derivative = derivative + (value * Decimal(weight))
@@ -256,32 +260,7 @@ class DerivativeSensor(RestoreEntity, SensorEntity):
             )
         )
 
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return self._name
-
     @property
     def native_value(self):
         """Return the state of the sensor."""
         return round(self._state, self._round_digits)
-
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit the value is expressed in."""
-        return self._unit_of_measurement
-
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
-    @property
-    def extra_state_attributes(self):
-        """Return the state attributes of the sensor."""
-        return {ATTR_SOURCE_ID: self._sensor_source_id}
-
-    @property
-    def icon(self):
-        """Return the icon to use in the frontend."""
-        return ICON