Improve type hints in integration (#78345)

homeassistant/components/integration/sensor.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 
 from decimal import Decimal, DecimalException
 import logging
+from typing import Final
 
 import voluptuous as vol
 
@@ -26,7 +27,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
@@ -45,11 +46,9 @@ from .const import (
     METHOD_TRAPEZOIDAL,
 )
 
-# mypy: allow-untyped-defs, no-check-untyped-defs
-
 _LOGGER = logging.getLogger(__name__)
 
-ATTR_SOURCE_ID = "source"
+ATTR_SOURCE_ID: Final = "source"
 
 # SI Metric prefixes
 UNIT_PREFIXES = {None: 1, "k": 10**3, "M": 10**6, "G": 10**9, "T": 10**12}
@@ -135,6 +134,9 @@ async def async_setup_platform(
 class IntegrationSensor(RestoreEntity, SensorEntity):
     """Representation of an integration sensor."""
 
+    _attr_state_class = SensorStateClass.TOTAL
+    _attr_should_poll = False
+
     def __init__(
         self,
         *,
@@ -155,13 +157,11 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
 
         self._attr_name = name if name is not None else f"{source_entity} integral"
         self._unit_template = f"{'' if unit_prefix is None else unit_prefix}{{}}"
-        self._unit_of_measurement = None
+        self._unit_of_measurement: str | None = None
         self._unit_prefix = UNIT_PREFIXES[unit_prefix]
         self._unit_time = UNIT_TIME[unit_time]
         self._unit_time_str = unit_time
-        self._attr_state_class = SensorStateClass.TOTAL
         self._attr_icon = "mdi:chart-histogram"
-        self._attr_should_poll = False
         self._attr_extra_state_attributes = {ATTR_SOURCE_ID: source_entity}
 
     def _unit(self, source_unit: str) -> str:
@@ -195,10 +195,10 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
                     )
 
         @callback
-        def calc_integration(event):
+        def calc_integration(event: Event) -> None:
             """Handle the sensor state changes."""
-            old_state = event.data.get("old_state")
-            new_state = event.data.get("new_state")
+            old_state: State | None = event.data.get("old_state")
+            new_state: State | None = event.data.get("new_state")
 
             if new_state is None or new_state.state in (
                 STATE_UNKNOWN,
@@ -237,7 +237,7 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
 
             try:
                 # integration as the Riemann integral of previous measures.
-                area = 0
+                area = Decimal(0)
                 elapsed_time = (
                     new_state.last_updated - old_state.last_updated
                 ).total_seconds()
@@ -277,13 +277,13 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
         )
 
     @property
-    def native_value(self):
+    def native_value(self) -> Decimal | None:
         """Return the state of the sensor."""
         if isinstance(self._state, Decimal):
             return round(self._state, self._round_digits)
         return self._state
 
     @property
-    def native_unit_of_measurement(self):
+    def native_unit_of_measurement(self) -> str | None:
         """Return the unit the value is expressed in."""
         return self._unit_of_measurement