Refactor event handlers

This commit is contained in:
Erik 2024-06-26 13:25:52 +02:00
parent 9cb8b5be54
commit 423f4c5bcf

View File

@ -178,7 +178,7 @@ _NAME_TO_INTEGRATION_METHOD: dict[str, type[_IntegrationMethod]] = {
class _IntegrationTrigger(Enum):
StateReport = "state_report"
StateEvent = "state_event"
TimeElapsed = "time_elapsed"
@ -338,7 +338,7 @@ class IntegrationSensor(RestoreSensor):
)
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.StateReport
self._last_integration_trigger = _IntegrationTrigger.StateEvent
self._attr_suggested_display_precision = round_digits or 2
def _calculate_unit(self, source_unit: str) -> str:
@ -428,9 +428,11 @@ class IntegrationSensor(RestoreSensor):
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_and_max_sub_interval
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)
@ -450,7 +452,7 @@ class IntegrationSensor(RestoreSensor):
self.async_on_remove(
self.hass.bus.async_listen(
EVENT_STATE_REPORTED,
handle_state_change,
handle_state_report,
event_filter=callback(
lambda event_data: event_data["entity_id"] == self._sensor_source_id
),
@ -459,8 +461,29 @@ class IntegrationSensor(RestoreSensor):
)
@callback
def _integrate_on_state_change_and_max_sub_interval(
self, event: Event[EventStateChangedData] | Event[EventStateReportedData]
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.
@ -468,21 +491,9 @@ class IntegrationSensor(RestoreSensor):
reschedules time based integration.
"""
self._cancel_max_sub_interval_exceeded_callback()
if event.event_type == EVENT_STATE_CHANGED:
if TYPE_CHECKING:
assert type(event.data) is EventStateChangedData
old_last_reported = None
old_state = event.data["old_state"]
else: # EVENT_STATE_REPORTED
if TYPE_CHECKING:
assert type(event.data) is EventStateReportedData
old_last_reported = event.data["old_last_reported"]
old_state = None
new_state = event.data["new_state"]
try:
self._integrate_on_state_change(old_last_reported, old_state, new_state)
self._last_integration_trigger = _IntegrationTrigger.StateReport
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
@ -491,21 +502,21 @@ class IntegrationSensor(RestoreSensor):
@callback
def _integrate_on_state_change_callback(
self, event: Event[EventStateChangedData] | Event[EventStateReportedData]
self, event: Event[EventStateChangedData]
) -> None:
"""Handle the sensor state changes."""
if event.event_type == EVENT_STATE_CHANGED:
if TYPE_CHECKING:
assert type(event.data) is EventStateChangedData
old_last_reported = None
old_state = event.data["old_state"]
else: # EVENT_STATE_REPORTED
if TYPE_CHECKING:
assert type(event.data) is EventStateReportedData
old_last_reported = event.data["old_last_reported"]
old_state = None
new_state = event.data["new_state"]
return self._integrate_on_state_change(old_last_reported, old_state, new_state)
"""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,
@ -546,7 +557,7 @@ class IntegrationSensor(RestoreSensor):
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.StateReport
if self._last_integration_trigger == _IntegrationTrigger.StateEvent
else (new_state.last_reported - self._last_integration_time).total_seconds()
)