Improve type hint in derivative sensor entity (#77038)

This commit is contained in:
epenet 2022-08-20 11:27:01 +02:00 committed by GitHub
parent 2c2e0cd4a0
commit 9ac01b8c9b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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)
# List of tuples with (timestamp_start, timestamp_end, derivative)
self._state_list: list[tuple[datetime, datetime, Decimal]] = []
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")
new_state = event.data.get("new_state")
old_state: State | None
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