Cleanup async_track_state_change and augment docstring (#37251)

* Cleanup async_track_state_change and augment docstrings.

Skip from_state and to_state matching in
async_track_state_change when they are None

Optimize the state change listener for the most
common use case: no to_state and from_state
matching.

* Update benchmark to be more realistic (previously we assumed only one entity was present in the whole instance)

* Add more tests to ensure behavior is preserved

* Ensure new behavior matches test

* remove MATCH_ALL from zone automation since its the default anyways

* Might as well use async_track_state_change_event instead since MATCH_ALL is removed

homeassistant/components/automation/zone.py

@@ -1,16 +1,10 @@
 """Offer zone automation rules."""
 import voluptuous as vol
 
-from homeassistant.const import (
-    CONF_ENTITY_ID,
-    CONF_EVENT,
-    CONF_PLATFORM,
-    CONF_ZONE,
-    MATCH_ALL,
-)
+from homeassistant.const import CONF_ENTITY_ID, CONF_EVENT, CONF_PLATFORM, CONF_ZONE
 from homeassistant.core import callback
 from homeassistant.helpers import condition, config_validation as cv, location
-from homeassistant.helpers.event import async_track_state_change
+from homeassistant.helpers.event import async_track_state_change_event
 
 # mypy: allow-untyped-defs, no-check-untyped-defs
 
@@ -37,8 +31,12 @@ async def async_attach_trigger(hass, config, action, automation_info):
     event = config.get(CONF_EVENT)
 
     @callback
-    def zone_automation_listener(entity, from_s, to_s):
+    def zone_automation_listener(zone_event):
         """Listen for state changes and calls action."""
+        entity = zone_event.data.get("entity_id")
+        from_s = zone_event.data.get("old_state")
+        to_s = zone_event.data.get("new_state")
+
         if (
             from_s
             and not location.has_location(from_s)
@@ -74,6 +72,4 @@ async def async_attach_trigger(hass, config, action, automation_info):
                 )
             )
 
-    return async_track_state_change(
-        hass, entity_id, zone_automation_listener, MATCH_ALL, MATCH_ALL
-    )
+    return async_track_state_change_event(hass, entity_id, zone_automation_listener)

homeassistant/helpers/event.py

@@ -72,10 +72,16 @@ def async_track_state_change(
 
     Returns a function that can be called to remove the listener.
 
+    If entity_ids are not MATCH_ALL along with from_state and to_state
+    being None, async_track_state_change_event should be used instead
+    as it is slightly faster.
+
     Must be run within the event loop.
     """
-    match_from_state = process_state_match(from_state)
-    match_to_state = process_state_match(to_state)
+    if from_state is not None:
+        match_from_state = process_state_match(from_state)
+    if to_state is not None:
+        match_to_state = process_state_match(to_state)
 
     # Ensure it is a lowercase list with entity ids we want to match on
     if entity_ids == MATCH_ALL:
@@ -88,21 +94,27 @@ def async_track_state_change(
     @callback
     def state_change_listener(event: Event) -> None:
         """Handle specific state changes."""
-        old_state = event.data.get("old_state")
-        if old_state is not None:
-            old_state = old_state.state
+        if from_state is not None:
+            old_state = event.data.get("old_state")
+            if old_state is not None:
+                old_state = old_state.state
 
-        new_state = event.data.get("new_state")
-        if new_state is not None:
-            new_state = new_state.state
+            if not match_from_state(old_state):
+                return
+        if to_state is not None:
+            new_state = event.data.get("new_state")
+            if new_state is not None:
+                new_state = new_state.state
 
-        if match_from_state(old_state) and match_to_state(new_state):
-            hass.async_run_job(
-                action,
-                event.data.get("entity_id"),
-                event.data.get("old_state"),
-                event.data.get("new_state"),
-            )
+            if not match_to_state(new_state):
+                return
+
+        hass.async_run_job(
+            action,
+            event.data.get("entity_id"),
+            event.data.get("old_state"),
+            event.data.get("new_state"),
+        )
 
     if entity_ids != MATCH_ALL:
         # If we have a list of entity ids we use
@@ -565,5 +577,5 @@ def process_state_match(
     if isinstance(parameter, str) or not hasattr(parameter, "__iter__"):
         return lambda state: state == parameter
 
-    parameter_tuple = tuple(parameter)
-    return lambda state: state in parameter_tuple
+    parameter_set = set(parameter)
+    return lambda state: state in parameter_set

homeassistant/scripts/benchmark/__init__.py

@@ -114,7 +114,7 @@ async def time_changed_helper(hass):
 
 @benchmark
 async def state_changed_helper(hass):
-    """Run a million events through state changed helper."""
+    """Run a million events through state changed helper with 1000 entities."""
     count = 0
     entity_id = "light.kitchen"
     event = asyncio.Event()
@@ -128,9 +128,48 @@ async def state_changed_helper(hass):
         if count == 10 ** 6:
             event.set()
 
-    hass.helpers.event.async_track_state_change(entity_id, listener, "off", "on")
+    for idx in range(1000):
+        hass.helpers.event.async_track_state_change(
+            f"{entity_id}{idx}", listener, "off", "on"
+        )
     event_data = {
-        "entity_id": entity_id,
+        "entity_id": f"{entity_id}0",
+        "old_state": core.State(entity_id, "off"),
+        "new_state": core.State(entity_id, "on"),
+    }
+
+    for _ in range(10 ** 6):
+        hass.bus.async_fire(EVENT_STATE_CHANGED, event_data)
+
+    start = timer()
+
+    await event.wait()
+
+    return timer() - start
+
+
+@benchmark
+async def state_changed_event_helper(hass):
+    """Run a million events through state changed event helper with 1000 entities."""
+    count = 0
+    entity_id = "light.kitchen"
+    event = asyncio.Event()
+
+    @core.callback
+    def listener(*args):
+        """Handle event."""
+        nonlocal count
+        count += 1
+
+        if count == 10 ** 6:
+            event.set()
+
+    hass.helpers.event.async_track_state_change_event(
+        [f"{entity_id}{idx}" for idx in range(1000)], listener
+    )
+
+    event_data = {
+        "entity_id": f"{entity_id}0",
         "old_state": core.State(entity_id, "off"),
         "new_state": core.State(entity_id, "on"),
     }

tests/helpers/test_event.py

@@ -81,6 +81,88 @@ async def test_track_point_in_time(hass):
     assert len(runs) == 2
 
 
+async def test_track_state_change_from_to_state_match(hass):
+    """Test track_state_change with from and to state matchers."""
+    from_and_to_state_runs = []
+    only_from_runs = []
+    only_to_runs = []
+    match_all_runs = []
+    no_to_from_specified_runs = []
+
+    def from_and_to_state_callback(entity_id, old_state, new_state):
+        from_and_to_state_runs.append(1)
+
+    def only_from_state_callback(entity_id, old_state, new_state):
+        only_from_runs.append(1)
+
+    def only_to_state_callback(entity_id, old_state, new_state):
+        only_to_runs.append(1)
+
+    def match_all_callback(entity_id, old_state, new_state):
+        match_all_runs.append(1)
+
+    def no_to_from_specified_callback(entity_id, old_state, new_state):
+        no_to_from_specified_runs.append(1)
+
+    async_track_state_change(
+        hass, "light.Bowl", from_and_to_state_callback, "on", "off"
+    )
+    async_track_state_change(hass, "light.Bowl", only_from_state_callback, "on", None)
+    async_track_state_change(hass, "light.Bowl", only_to_state_callback, None, "off")
+    async_track_state_change(
+        hass, "light.Bowl", match_all_callback, MATCH_ALL, MATCH_ALL
+    )
+    async_track_state_change(hass, "light.Bowl", no_to_from_specified_callback)
+
+    hass.states.async_set("light.Bowl", "on")
+    await hass.async_block_till_done()
+    assert len(from_and_to_state_runs) == 0
+    assert len(only_from_runs) == 0
+    assert len(only_to_runs) == 0
+    assert len(match_all_runs) == 1
+    assert len(no_to_from_specified_runs) == 1
+
+    hass.states.async_set("light.Bowl", "off")
+    await hass.async_block_till_done()
+    assert len(from_and_to_state_runs) == 1
+    assert len(only_from_runs) == 1
+    assert len(only_to_runs) == 1
+    assert len(match_all_runs) == 2
+    assert len(no_to_from_specified_runs) == 2
+
+    hass.states.async_set("light.Bowl", "on")
+    await hass.async_block_till_done()
+    assert len(from_and_to_state_runs) == 1
+    assert len(only_from_runs) == 1
+    assert len(only_to_runs) == 1
+    assert len(match_all_runs) == 3
+    assert len(no_to_from_specified_runs) == 3
+
+    hass.states.async_set("light.Bowl", "on")
+    await hass.async_block_till_done()
+    assert len(from_and_to_state_runs) == 1
+    assert len(only_from_runs) == 1
+    assert len(only_to_runs) == 1
+    assert len(match_all_runs) == 3
+    assert len(no_to_from_specified_runs) == 3
+
+    hass.states.async_set("light.Bowl", "off")
+    await hass.async_block_till_done()
+    assert len(from_and_to_state_runs) == 2
+    assert len(only_from_runs) == 2
+    assert len(only_to_runs) == 2
+    assert len(match_all_runs) == 4
+    assert len(no_to_from_specified_runs) == 4
+
+    hass.states.async_set("light.Bowl", "off")
+    await hass.async_block_till_done()
+    assert len(from_and_to_state_runs) == 2
+    assert len(only_from_runs) == 2
+    assert len(only_to_runs) == 2
+    assert len(match_all_runs) == 4
+    assert len(no_to_from_specified_runs) == 4
+
+
 async def test_track_state_change(hass):
     """Test track_state_change."""
     # 2 lists to track how often our callbacks get called
@@ -91,12 +173,14 @@ async def test_track_state_change(hass):
     def specific_run_callback(entity_id, old_state, new_state):
         specific_runs.append(1)
 
+    # This is the rare use case
     async_track_state_change(hass, "light.Bowl", specific_run_callback, "on", "off")
 
     @ha.callback
     def wildcard_run_callback(entity_id, old_state, new_state):
         wildcard_runs.append((old_state, new_state))
 
+    # This is the most common use case
     async_track_state_change(hass, "light.Bowl", wildcard_run_callback)
 
     async def wildercard_run_callback(entity_id, old_state, new_state):