diff --git a/esphome/components/select/__init__.py b/esphome/components/select/__init__.py index 7e1dc8e6b42..8e0816d75d5 100644 --- a/esphome/components/select/__init__.py +++ b/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) diff --git a/esphome/core/entity_helpers.py b/esphome/core/entity_helpers.py index dc526c4e970..2722567547d 100644 --- a/esphome/core/entity_helpers.py +++ b/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: diff --git a/tests/unit_tests/core/test_entity_helpers.py b/tests/unit_tests/core/test_entity_helpers.py index 90e867ac9e4..4d42e12e541 100644 --- a/tests/unit_tests/core/test_entity_helpers.py +++ b/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]]