"""Numeric integration of data coming from a source sensor over time."""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from decimal import Decimal, InvalidOperation
from enum import Enum
import logging
from typing import TYPE_CHECKING, Any, Final, Self
import voluptuous as vol
from homeassistant.components.sensor import (
DEVICE_CLASS_UNITS,
PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA,
RestoreSensor,
SensorDeviceClass,
SensorExtraStoredData,
SensorStateClass,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import (
ATTR_DEVICE_CLASS,
ATTR_UNIT_OF_MEASUREMENT,
CONF_METHOD,
CONF_NAME,
CONF_UNIQUE_ID,
STATE_UNAVAILABLE,
UnitOfTime,
)
from homeassistant.core import (
CALLBACK_TYPE,
Event,
EventStateChangedData,
EventStateReportedData,
HomeAssistant,
State,
callback,
)
from homeassistant.helpers import config_validation as cv, entity_registry as er
from homeassistant.helpers.device import async_device_info_to_link_from_entity
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import (
async_call_later,
async_track_state_change_event,
async_track_state_report_event,
)
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
from .const import (
CONF_MAX_SUB_INTERVAL,
CONF_ROUND_DIGITS,
CONF_SOURCE_SENSOR,
CONF_UNIT_OF_MEASUREMENT,
CONF_UNIT_PREFIX,
CONF_UNIT_TIME,
INTEGRATION_METHODS,
METHOD_LEFT,
METHOD_RIGHT,
METHOD_TRAPEZOIDAL,
)
_LOGGER = logging.getLogger(__name__)
ATTR_SOURCE_ID: Final = "source"
# SI Metric prefixes
UNIT_PREFIXES = {None: 1, "k": 10**3, "M": 10**6, "G": 10**9, "T": 10**12}
# SI Time prefixes
UNIT_TIME = {
UnitOfTime.SECONDS: 1,
UnitOfTime.MINUTES: 60,
UnitOfTime.HOURS: 60 * 60,
UnitOfTime.DAYS: 24 * 60 * 60,
}
DEVICE_CLASS_MAP = {
SensorDeviceClass.POWER: SensorDeviceClass.ENERGY,
}
DEFAULT_ROUND = 3
PLATFORM_SCHEMA = vol.All(
cv.removed(CONF_UNIT_OF_MEASUREMENT),
SENSOR_PLATFORM_SCHEMA.extend(
{
vol.Optional(CONF_NAME): cv.string,
vol.Optional(CONF_UNIQUE_ID): cv.string,
vol.Required(CONF_SOURCE_SENSOR): cv.entity_id,
vol.Optional(CONF_ROUND_DIGITS, default=DEFAULT_ROUND): vol.Any(
None, vol.Coerce(int)
),
vol.Optional(CONF_UNIT_PREFIX): vol.In(UNIT_PREFIXES),
vol.Optional(CONF_UNIT_TIME, default=UnitOfTime.HOURS): vol.In(UNIT_TIME),
vol.Remove(CONF_UNIT_OF_MEASUREMENT): cv.string,
vol.Optional(CONF_MAX_SUB_INTERVAL): cv.positive_time_period,
vol.Optional(CONF_METHOD, default=METHOD_TRAPEZOIDAL): vol.In(
INTEGRATION_METHODS
),
}
),
)
class _IntegrationMethod(ABC):
@staticmethod
def from_name(method_name: str) -> _IntegrationMethod:
return _NAME_TO_INTEGRATION_METHOD[method_name]()
@abstractmethod
def validate_states(self, left: str, right: str) -> tuple[Decimal, Decimal] | None:
"""Check state requirements for integration."""
@abstractmethod
def calculate_area_with_two_states(
self, elapsed_time: Decimal, left: Decimal, right: Decimal
) -> Decimal:
"""Calculate area given two states."""
def calculate_area_with_one_state(
self, elapsed_time: Decimal, constant_state: Decimal
) -> Decimal:
return constant_state * elapsed_time
class _Trapezoidal(_IntegrationMethod):
def calculate_area_with_two_states(
self, elapsed_time: Decimal, left: Decimal, right: Decimal
) -> Decimal:
return elapsed_time * (left + right) / 2
def validate_states(self, left: str, right: str) -> tuple[Decimal, Decimal] | None:
if (left_dec := _decimal_state(left)) is None or (
right_dec := _decimal_state(right)
) is None:
return None
return (left_dec, right_dec)
class _Left(_IntegrationMethod):
def calculate_area_with_two_states(
self, elapsed_time: Decimal, left: Decimal, right: Decimal
) -> Decimal:
return self.calculate_area_with_one_state(elapsed_time, left)
def validate_states(self, left: str, right: str) -> tuple[Decimal, Decimal] | None:
if (left_dec := _decimal_state(left)) is None:
return None
return (left_dec, left_dec)
class _Right(_IntegrationMethod):
def calculate_area_with_two_states(
self, elapsed_time: Decimal, left: Decimal, right: Decimal
) -> Decimal:
return self.calculate_area_with_one_state(elapsed_time, right)
def validate_states(self, left: str, right: str) -> tuple[Decimal, Decimal] | None:
if (right_dec := _decimal_state(right)) is None:
return None
return (right_dec, right_dec)
def _decimal_state(state: str) -> Decimal | None:
try:
return Decimal(state)
except (InvalidOperation, TypeError):
return None
_NAME_TO_INTEGRATION_METHOD: dict[str, type[_IntegrationMethod]] = {
METHOD_LEFT: _Left,
METHOD_RIGHT: _Right,
METHOD_TRAPEZOIDAL: _Trapezoidal,
}
class _IntegrationTrigger(Enum):
StateEvent = "state_event"
TimeElapsed = "time_elapsed"
@dataclass
class IntegrationSensorExtraStoredData(SensorExtraStoredData):
"""Object to hold extra stored data."""
source_entity: str | None
last_valid_state: Decimal | None
def as_dict(self) -> dict[str, Any]:
"""Return a dict representation of the utility sensor data."""
data = super().as_dict()
data["source_entity"] = self.source_entity
data["last_valid_state"] = (
str(self.last_valid_state) if self.last_valid_state else None
)
return data
@classmethod
def from_dict(cls, restored: dict[str, Any]) -> Self | None:
"""Initialize a stored sensor state from a dict."""
extra = SensorExtraStoredData.from_dict(restored)
if extra is None:
return None
source_entity = restored.get(ATTR_SOURCE_ID)
try:
last_valid_state = (
Decimal(str(restored.get("last_valid_state")))
if restored.get("last_valid_state")
else None
)
except InvalidOperation:
# last_period is corrupted
_LOGGER.error("Could not use last_valid_state")
return None
if last_valid_state is None:
return None
return cls(
extra.native_value,
extra.native_unit_of_measurement,
source_entity,
last_valid_state,
)
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Initialize Integration - Riemann sum integral config entry."""
registry = er.async_get(hass)
# Validate + resolve entity registry id to entity_id
source_entity_id = er.async_validate_entity_id(
registry, config_entry.options[CONF_SOURCE_SENSOR]
)
device_info = async_device_info_to_link_from_entity(
hass,
source_entity_id,
)
if (unit_prefix := config_entry.options.get(CONF_UNIT_PREFIX)) == "none":
# Before we had support for optional selectors, "none" was used for selecting nothing
unit_prefix = None
if max_sub_interval_dict := config_entry.options.get(CONF_MAX_SUB_INTERVAL, None):
max_sub_interval = cv.time_period(max_sub_interval_dict)
else:
max_sub_interval = None
round_digits = config_entry.options.get(CONF_ROUND_DIGITS)
if round_digits:
round_digits = int(round_digits)
integral = IntegrationSensor(
integration_method=config_entry.options[CONF_METHOD],
name=config_entry.title,
round_digits=round_digits,
source_entity=source_entity_id,
unique_id=config_entry.entry_id,
unit_prefix=unit_prefix,
unit_time=config_entry.options[CONF_UNIT_TIME],
device_info=device_info,
max_sub_interval=max_sub_interval,
)
async_add_entities([integral])
async def async_setup_platform(
hass: HomeAssistant,
config: ConfigType,
async_add_entities: AddEntitiesCallback,
discovery_info: DiscoveryInfoType | None = None,
) -> None:
"""Set up the integration sensor."""
integral = IntegrationSensor(
integration_method=config[CONF_METHOD],
name=config.get(CONF_NAME),
round_digits=config.get(CONF_ROUND_DIGITS),
source_entity=config[CONF_SOURCE_SENSOR],
unique_id=config.get(CONF_UNIQUE_ID),
unit_prefix=config.get(CONF_UNIT_PREFIX),
unit_time=config[CONF_UNIT_TIME],
max_sub_interval=config.get(CONF_MAX_SUB_INTERVAL),
)
async_add_entities([integral])
class IntegrationSensor(RestoreSensor):
"""Representation of an integration sensor."""
_attr_state_class = SensorStateClass.TOTAL
_attr_should_poll = False
def __init__(
self,
*,
integration_method: str,
name: str | None,
round_digits: int | None,
source_entity: str,
unique_id: str | None,
unit_prefix: str | None,
unit_time: UnitOfTime,
max_sub_interval: timedelta | None,
device_info: DeviceInfo | None = None,
) -> None:
"""Initialize the integration sensor."""
self._attr_unique_id = unique_id
self._sensor_source_id = source_entity
self._round_digits = round_digits
self._state: Decimal | None = None
self._method = _IntegrationMethod.from_name(integration_method)
self._attr_name = name if name is not None else f"{source_entity} integral"
self._unit_prefix_string = "" if unit_prefix is None else unit_prefix
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_icon = "mdi:chart-histogram"
self._source_entity: str = source_entity
self._last_valid_state: Decimal | None = None
self._attr_device_info = device_info
self._max_sub_interval: timedelta | None = (
None # disable time based integration
if max_sub_interval is None or max_sub_interval.total_seconds() == 0
else max_sub_interval
)
self._max_sub_interval_exceeded_callback: CALLBACK_TYPE = lambda *args: None
self._last_integration_time: datetime = datetime.now(tz=UTC)
self._last_integration_trigger = _IntegrationTrigger.StateEvent
self._attr_suggested_display_precision = round_digits or 2
def _calculate_unit(self, source_unit: str) -> str:
"""Multiply source_unit with time unit of the integral.
Possibly cancelling out a time unit in the denominator of the source_unit.
Note that this is a heuristic string manipulation method and might not
transform all source units in a sensible way.
Examples:
- Speed to distance: 'km/h' and 'h' will be transformed to 'km'
- Power to energy: 'W' and 'h' will be transformed to 'Wh'
"""
unit_time = self._unit_time_str
if source_unit.endswith(f"/{unit_time}"):
integral_unit = source_unit[0 : (-(1 + len(unit_time)))]
else:
integral_unit = f"{source_unit}{unit_time}"
return f"{self._unit_prefix_string}{integral_unit}"
def _calculate_device_class(
self,
source_device_class: SensorDeviceClass | None,
unit_of_measurement: str | None,
) -> SensorDeviceClass | None:
"""Deduce device class if possible from source device class and target unit."""
if source_device_class is None:
return None
if (device_class := DEVICE_CLASS_MAP.get(source_device_class)) is None:
return None
if unit_of_measurement not in DEVICE_CLASS_UNITS.get(device_class, set()):
return None
return device_class
def _derive_and_set_attributes_from_state(self, source_state: State) -> None:
source_unit = source_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
if source_unit is not None:
self._unit_of_measurement = self._calculate_unit(source_unit)
else:
# If the source has no defined unit we cannot derive a unit for the integral
self._unit_of_measurement = None
self._attr_device_class = self._calculate_device_class(
source_state.attributes.get(ATTR_DEVICE_CLASS), self.unit_of_measurement
)
if self._attr_device_class:
self._attr_icon = None # Remove this sensors icon default and allow to fallback to the device class default
else:
self._attr_icon = "mdi:chart-histogram"
def _update_integral(self, area: Decimal) -> None:
area_scaled = area / (self._unit_prefix * self._unit_time)
if isinstance(self._state, Decimal):
self._state += area_scaled
else:
self._state = area_scaled
_LOGGER.debug(
"area = %s, area_scaled = %s new state = %s", area, area_scaled, self._state
)
self._last_valid_state = self._state
async def async_added_to_hass(self) -> None:
"""Handle entity which will be added."""
await super().async_added_to_hass()
if (last_sensor_data := await self.async_get_last_sensor_data()) is not None:
self._state = (
Decimal(str(last_sensor_data.native_value))
if last_sensor_data.native_value
else last_sensor_data.last_valid_state
)
self._attr_native_value = last_sensor_data.native_value
self._unit_of_measurement = last_sensor_data.native_unit_of_measurement
self._last_valid_state = last_sensor_data.last_valid_state
_LOGGER.debug(
"Restored state %s and last_valid_state %s",
self._state,
self._last_valid_state,
)
if self._max_sub_interval is not None:
source_state = self.hass.states.get(self._sensor_source_id)
self._schedule_max_sub_interval_exceeded_if_state_is_numeric(source_state)
self.async_on_remove(self._cancel_max_sub_interval_exceeded_callback)
handle_state_change = self._integrate_on_state_change_with_max_sub_interval
handle_state_report = self._integrate_on_state_report_with_max_sub_interval
else:
handle_state_change = self._integrate_on_state_change_callback
handle_state_report = self._integrate_on_state_report_callback
if (
state := self.hass.states.get(self._source_entity)
) and state.state != STATE_UNAVAILABLE:
self._derive_and_set_attributes_from_state(state)
self.async_on_remove(
async_track_state_change_event(
self.hass,
self._sensor_source_id,
handle_state_change,
)
)
self.async_on_remove(
async_track_state_report_event(
self.hass,
self._sensor_source_id,
handle_state_report,
)
)
@callback
def _integrate_on_state_change_with_max_sub_interval(
self, event: Event[EventStateChangedData]
) -> None:
"""Handle sensor state update when sub interval is configured."""
self._integrate_on_state_update_with_max_sub_interval(
None, event.data["old_state"], event.data["new_state"]
)
@callback
def _integrate_on_state_report_with_max_sub_interval(
self, event: Event[EventStateReportedData]
) -> None:
"""Handle sensor state report when sub interval is configured."""
self._integrate_on_state_update_with_max_sub_interval(
event.data["old_last_reported"], None, event.data["new_state"]
)
@callback
def _integrate_on_state_update_with_max_sub_interval(
self,
old_last_reported: datetime | None,
old_state: State | None,
new_state: State | None,
) -> None:
"""Integrate based on state change and time.
Next to doing the integration based on state change this method cancels and
reschedules time based integration.
"""
self._cancel_max_sub_interval_exceeded_callback()
try:
self._integrate_on_state_change(old_last_reported, old_state, new_state)
self._last_integration_trigger = _IntegrationTrigger.StateEvent
self._last_integration_time = datetime.now(tz=UTC)
finally:
# When max_sub_interval exceeds without state change the source is assumed
# constant with the last known state (new_state).
self._schedule_max_sub_interval_exceeded_if_state_is_numeric(new_state)
@callback
def _integrate_on_state_change_callback(
self, event: Event[EventStateChangedData]
) -> None:
"""Handle sensor state change."""
return self._integrate_on_state_change(
None, event.data["old_state"], event.data["new_state"]
)
@callback
def _integrate_on_state_report_callback(
self, event: Event[EventStateReportedData]
) -> None:
"""Handle sensor state report."""
return self._integrate_on_state_change(
event.data["old_last_reported"], None, event.data["new_state"]
)
def _integrate_on_state_change(
self,
old_last_reported: datetime | None,
old_state: State | None,
new_state: State | None,
) -> None:
if new_state is None:
return
if new_state.state == STATE_UNAVAILABLE:
self._attr_available = False
self.async_write_ha_state()
return
if old_state:
# state has changed, we recover old_state from the event
old_state_state = old_state.state
old_last_reported = old_state.last_reported
else:
# event state reported without any state change
old_state_state = new_state.state
self._attr_available = True
self._derive_and_set_attributes_from_state(new_state)
if old_last_reported is None and old_state is None:
self.async_write_ha_state()
return
if not (
states := self._method.validate_states(old_state_state, new_state.state)
):
self.async_write_ha_state()
return
if TYPE_CHECKING:
assert old_last_reported is not None
elapsed_seconds = Decimal(
(new_state.last_reported - old_last_reported).total_seconds()
if self._last_integration_trigger == _IntegrationTrigger.StateEvent
else (new_state.last_reported - self._last_integration_time).total_seconds()
)
area = self._method.calculate_area_with_two_states(elapsed_seconds, *states)
self._update_integral(area)
self.async_write_ha_state()
def _schedule_max_sub_interval_exceeded_if_state_is_numeric(
self, source_state: State | None
) -> None:
"""Schedule possible integration using the source state and max_sub_interval.
The callback reference is stored for possible cancellation if the source state
reports a change before max_sub_interval has passed.
If the callback is executed, meaning there was no state change reported, the
source_state is assumed constant and integration is done using its value.
"""
if (
self._max_sub_interval is not None
and source_state is not None
and (source_state_dec := _decimal_state(source_state.state))
):
@callback
def _integrate_on_max_sub_interval_exceeded_callback(now: datetime) -> None:
"""Integrate based on time and reschedule."""
elapsed_seconds = Decimal(
(now - self._last_integration_time).total_seconds()
)
self._derive_and_set_attributes_from_state(source_state)
area = self._method.calculate_area_with_one_state(
elapsed_seconds, source_state_dec
)
self._update_integral(area)
self.async_write_ha_state()
self._last_integration_time = datetime.now(tz=UTC)
self._last_integration_trigger = _IntegrationTrigger.TimeElapsed
self._schedule_max_sub_interval_exceeded_if_state_is_numeric(
source_state
)
self._max_sub_interval_exceeded_callback = async_call_later(
self.hass,
self._max_sub_interval,
_integrate_on_max_sub_interval_exceeded_callback,
)
def _cancel_max_sub_interval_exceeded_callback(self) -> None:
self._max_sub_interval_exceeded_callback()
@property
def native_value(self) -> Decimal | None:
"""Return the state of the sensor."""
if isinstance(self._state, Decimal) and self._round_digits:
return round(self._state, self._round_digits)
return self._state
@property
def native_unit_of_measurement(self) -> str | None:
"""Return the unit the value is expressed in."""
return self._unit_of_measurement
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the sensor."""
return {
ATTR_SOURCE_ID: self._source_entity,
}
@property
def extra_restore_state_data(self) -> IntegrationSensorExtraStoredData:
"""Return sensor specific state data to be restored."""
return IntegrationSensorExtraStoredData(
self.native_value,
self.native_unit_of_measurement,
self._source_entity,
self._last_valid_state,
)
async def async_get_last_sensor_data(
self,
) -> IntegrationSensorExtraStoredData | None:
"""Restore Utility Meter Sensor Extra Stored Data."""
if (restored_last_extra_data := await self.async_get_last_extra_data()) is None:
return None
return IntegrationSensorExtraStoredData.from_dict(
restored_last_extra_data.as_dict()
)