Merge branch 'apply-entities' into apply-entities-number

esphome/components/select/__init__.py

@@ -130,13 +130,14 @@ APPLY_ENTITY = ApplyEntity(
     ApplySelect,
     ((SelectPtr, "entity"), (cg.size_t, "index")),
     "entity->publish_state(index);",
+    (("option", "const char *option = entity->option_at(index);"),),
 )
 
 
 def sub_selects(
     config: ConfigType, *, parent: MockObj | ID | None = None
 ) -> SubEntities:
-    """Return a SubEntities bound to new_select; ``apply`` statements can use ``entity`` and ``index``."""
+    """Return a SubEntities bound to new_select; ``apply`` statements can use ``entity``, ``index`` and ``option``."""
     return SubEntities(new_select, config, parent, APPLY_ENTITY)
 
 

esphome/core/entity_helpers.py

@@ -2,6 +2,7 @@ from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
 import functools
 import logging
+import re
 from typing import Any, Literal
 
 from esphome.automation import apply_static_function, parent_ref
@@ -474,11 +475,15 @@ ApplyPublish = Literal["before", "after"] | None
 
 @dataclass(frozen=True, slots=True)
 class ApplyEntity:
-    """A domain's ``Apply*`` class, its generated function's parameters and its publish statement."""
+    """A domain's ``Apply*`` class, its generated function's parameters and its publish statement.
+
+    ``locals`` are (name, declaration) pairs, emitted only when the body uses the name.
+    """
 
     cls: MockObjClass
     params: tuple[tuple[SafeExpType, str], ...]
     publish: str | None = None
+    locals: tuple[tuple[str, str], ...] = ()
 
 
 _APPLY_DOMAIN = "entity_apply"
@@ -562,6 +567,12 @@ class SubEntities:
             body.insert(0, apply_entity.publish)
         elif publish == "after":
             body.append(apply_entity.publish)
+        text = "\n".join(body)
+        body[:0] = [
+            decl
+            for name, decl in apply_entity.locals
+            if re.search(rf"\b{name}\b", text)
+        ]
         functions: dict[str, SafeExpType] = CORE.data.setdefault(_APPLY_DOMAIN, {})
         key = "\n".join([*(f"{t} {n}" for t, n in apply_entity.params), *body])
         if (fn := functions.get(key)) is None:

tests/unit_tests/core/test_entity_helpers.py

@@ -1386,6 +1386,46 @@ async def test_sub_entities_apply_generates_function(
     assert setup_test_environment == ["hub->set_bt_switch(bt_sw)"]
 
 
+_SELECT_APPLY = entity_helpers.ApplyEntity(
+    cg.esphome_ns.namespace("select").class_("ApplySelect"),
+    ((RawExpression("select::Select *"), "entity"), (RawExpression("size_t"), "index")),
+    "entity->publish_state(index);",
+    (("option", "const char *option = entity->option_at(index);"),),
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("apply", "body"),
+    [
+        (
+            "set_mode(option)",
+            (
+                "const char *option = entity->option_at(index);\n"
+                "::hub->set_mode(option);\nentity->publish_state(index);"
+            ),
+        ),
+        (
+            "set_mode(index, options_)",
+            "::hub->set_mode(index, options_);\nentity->publish_state(index);",
+        ),
+    ],
+)
+async def test_sub_entities_apply_locals(
+    setup_test_environment: list[str],
+    apply_functions: list[tuple[str, str, str]],
+    apply: str,
+    body: str,
+) -> None:
+    """A local is declared only when the body names it as a whole identifier."""
+    conf = _apply_conf("m_sel", _SELECT_APPLY.cls)
+    sub = _apply_sub([], {"m": conf}, _SELECT_APPLY)
+    await sub("m", None, apply=apply, publish="after")
+    assert apply_functions == [
+        ("esphome__m_sel__fn", "select::Select * entity, size_t index", body)
+    ]
+
+
 @pytest.mark.asyncio
 async def test_sub_entities_apply_shares_equal_functions(
     setup_test_environment: list[str], apply_functions: list[tuple[str, str, str]]