Avoid enumerating the whole state machine to find zone entities (#94866)

This commit is contained in:
J. Nick Koston 2023-06-20 14:48:28 +01:00 committed by GitHub
parent 30e8f806c1
commit b51dcb600e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -99,6 +99,8 @@ STORAGE_VERSION = 1
ENTITY_ID_SORTER = attrgetter("entity_id")
ZONE_ENTITY_IDS = "zone_entity_ids"
@bind_hass
def async_active_zone(
@ -111,9 +113,15 @@ def async_active_zone(
# Sort entity IDs so that we are deterministic if equal distance to 2 zones
min_dist = None
closest = None
for zone in sorted(hass.states.async_all(DOMAIN), key=ENTITY_ID_SORTER):
if zone.state == STATE_UNAVAILABLE or zone.attributes.get(ATTR_PASSIVE):
# This can be called before async_setup by device tracker
zone_entity_ids: list[str] = hass.data.get(ZONE_ENTITY_IDS, [])
for entity_id in zone_entity_ids:
zone = hass.states.get(entity_id)
if (
not zone
or zone.state == STATE_UNAVAILABLE
or zone.attributes.get(ATTR_PASSIVE)
):
continue
zone_dist = distance(
@ -141,6 +149,27 @@ def async_active_zone(
return closest
@callback
def async_setup_track_zone_entity_ids(hass: HomeAssistant) -> None:
"""Set up track of entity IDs for zones."""
zone_entity_ids: list[str] = hass.states.async_entity_ids(DOMAIN)
hass.data[ZONE_ENTITY_IDS] = zone_entity_ids
@callback
def _async_add_zone_entity_id(event_: Event) -> None:
"""Add zone entity ID."""
zone_entity_ids.append(event_.data[ATTR_ENTITY_ID])
zone_entity_ids.sort()
@callback
def _async_remove_zone_entity_id(event_: Event) -> None:
"""Remove zone entity ID."""
zone_entity_ids.remove(event_.data[ATTR_ENTITY_ID])
event.async_track_state_added_domain(hass, DOMAIN, _async_add_zone_entity_id)
event.async_track_state_removed_domain(hass, DOMAIN, _async_remove_zone_entity_id)
def in_zone(zone: State, latitude: float, longitude: float, radius: float = 0) -> bool:
"""Test if given latitude, longitude is in given zone.
@ -184,6 +213,8 @@ class ZoneStorageCollection(collection.DictStorageCollection):
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up configured zones as well as Home Assistant zone if necessary."""
async_setup_track_zone_entity_ids(hass)
component = entity_component.EntityComponent[Zone](_LOGGER, DOMAIN, hass)
id_manager = collection.IDManager()