[core] Share the apply function naming with automations and check the Apply class at codegen

esphome/automation.py

@@ -426,12 +426,20 @@ def _apply_function(
 ) -> MockObj:
     """Emit the generated function and declare ``id_`` as the ``ApplyAction`` or
     ``ApplyCondition`` templated on it, so ``play()`` calls it directly."""
-    fn = cg.static_function(
-        f"esphome__{id_.id}__fn", return_type, lambda_args, statements
-    )
+    fn = apply_static_function(id_, return_type, lambda_args, statements)
     return cg.new_Pvariable(id_, cg.TemplateArguments(fn, *template_arg))
 
 
+def apply_static_function(
+    id_: ID,
+    return_type: SafeExpType,
+    params: TemplateArgsType,
+    statements: list[str],
+) -> MockObj:
+    """Emit the static function generated for ``id_`` and return an expression naming it."""
+    return cg.static_function(f"esphome__{id_.id}__fn", return_type, params, statements)
+
+
 async def _render_values(
     name: str,
     target: str,

esphome/components/button/__init__.py

@@ -118,7 +118,7 @@ async def new_button(config: ConfigType, *args: SafeExpType) -> MockObj:
     return var
 
 
-_APPLY_ENTITY = ApplyEntity(((ButtonPtr, "entity"),))
+_APPLY_ENTITY = ApplyEntity(ApplyButton, ((ButtonPtr, "entity"),))
 
 
 def sub_buttons(

esphome/components/number/__init__.py

@@ -331,7 +331,9 @@ async def new_number(config, *args, min_value: float, max_value: float, step: fl
 
 
 _APPLY_ENTITY = ApplyEntity(
-    ((NumberPtr, "entity"), (cg.float_, "value")), "entity->publish_state(value);"
+    ApplyNumber,
+    ((NumberPtr, "entity"), (cg.float_, "value")),
+    "entity->publish_state(value);",
 )
 
 

esphome/components/select/__init__.py

@@ -127,7 +127,9 @@ async def new_select(config, *args, options: list[str]):
 
 
 _APPLY_ENTITY = ApplyEntity(
-    ((SelectPtr, "entity"), (cg.size_t, "index")), "entity->publish_state(index);"
+    ApplySelect,
+    ((SelectPtr, "entity"), (cg.size_t, "index")),
+    "entity->publish_state(index);",
 )
 
 

esphome/components/switch/__init__.py

@@ -179,7 +179,9 @@ async def new_switch(config, *args):
 
 
 _APPLY_ENTITY = ApplyEntity(
-    ((SwitchPtr, "entity"), (cg.bool_, "state")), "entity->publish_state(state);"
+    ApplySwitch,
+    ((SwitchPtr, "entity"), (cg.bool_, "state")),
+    "entity->publish_state(state);",
 )
 
 

esphome/core/entity_helpers.py

@@ -2,8 +2,9 @@ from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
 import functools
 import logging
-from typing import Any
+from typing import Any, Literal
 
+from esphome.automation import apply_static_function, parent_ref
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.const import (
@@ -23,7 +24,7 @@ from esphome.core.config import (
     ICON_MAX_LENGTH,
     UNIT_OF_MEASUREMENT_MAX_LENGTH,
 )
-from esphome.cpp_generator import MockObj, RawStatement, add, get_variable
+from esphome.cpp_generator import MockObj, MockObjClass, RawStatement, add, get_variable
 from esphome.cpp_types import App
 import esphome.final_validate as fv
 from esphome.helpers import (
@@ -468,10 +469,14 @@ async def _setup_entity_impl(var: MockObj, config: ConfigType, platform: str) ->
     config[_KEY_ICON_IDX] = icon_idx
 
 
+ApplyPublish = Literal["before", "after"] | None
+
+
 @dataclass(frozen=True, slots=True)
 class ApplyEntity:
-    """A domain's generated control function: its parameters and how it publishes the new value."""
+    """A domain's ``Apply*`` class, its generated function's parameters and its publish statement."""
 
+    cls: MockObjClass
     params: tuple[tuple[SafeExpType, str], ...]
     publish: str | None = None
 
@@ -497,8 +502,8 @@ def _apply_entity_function(
     functions = _get_apply_data().apply_functions
     key = "\n".join([*(f"{type_} {name}" for type_, name in params), *body])
     if (fn := functions.get(key)) is None:
-        fn = functions[key] = cg.static_function(
-            f"esphome__{entity_id.id}__fn", cg.void, list(params), body
+        fn = functions[key] = apply_static_function(
+            entity_id, cg.void, list(params), body
         )
     return fn
 
@@ -507,12 +512,11 @@ def _apply_entity_function(
 class SubEntities:
     """Create the child entity configured under key, set parent (if bound) and pass it to setter.
 
-    ``setter`` may be None when the hub does not keep the entity.
-
-    Extra arguments go to new_entity. Returns None only when key is absent, so the result can be
-    used directly as a condition. With ``apply`` the entity is the domain's ``Apply*`` class and
-    its control runs ``parent->apply`` (see ``ApplyEntity.params`` for the names it can use);
-    ``publish`` is ``"before"`` or ``"after"`` that call, or None to not publish.
+    ``setter`` may be None when the hub does not keep the entity. Extra arguments go to
+    new_entity. Returns None only when key is absent, so the result can be used directly as a
+    condition. With ``apply`` the entity is the domain's ``Apply*`` class and its control runs
+    ``parent->apply`` (see ``ApplyEntity.params`` for the names it can use); ``publish`` is
+    ``"before"`` or ``"after"`` that call, or None to not publish.
     """
 
     new_entity: Callable[..., Awaitable[MockObj]]
@@ -524,16 +528,12 @@ class SubEntities:
         self,
         key: str,
         setter: Callable[[MockObj], Expression] | None,
-        *args: SafeExpType,
-        apply: str | None = None,
-        publish: str | None = None,
+        *args: Any,
         **kwargs: Any,
     ) -> MockObj | None:
         if (conf := self.config.get(key)) is None:
             return None
-        return await self.create(
-            conf, setter, *args, apply=apply, publish=publish, **kwargs
-        )
+        return await self.create(conf, setter, *args, **kwargs)
 
     async def create(
         self,
@@ -541,39 +541,43 @@ class SubEntities:
         setter: Callable[[MockObj], Expression] | None,
         *args: SafeExpType,
         apply: str | None = None,
-        publish: str | None = None,
+        publish: ApplyPublish = None,
         **kwargs: Any,
     ) -> MockObj:
         """Create the entity from ``conf`` itself, for a platform whose config is one entity."""
+        entity_id: ID = conf[CONF_ID]
+        is_apply_class = self.apply_entity is not None and str(entity_id.type) == str(
+            self.apply_entity.cls
+        )
+        if apply is None and is_apply_class:
+            raise ValueError(f"{entity_id}: {entity_id.type} needs apply=")
+        if apply is not None and not is_apply_class:
+            raise ValueError(
+                f"{entity_id}: apply= needs the domain's Apply class, not {entity_id.type}"
+            )
+        if apply is not None and self.parent is None:
+            raise ValueError(f"{entity_id}: apply= needs a parent")
+        if publish is not None and (apply is None or self.apply_entity.publish is None):
+            raise ValueError(f"{entity_id}: publish= needs apply= on a stateful domain")
+
         if apply is None:
-            if publish is not None:
-                raise ValueError(f"{conf[CONF_ID]}: publish needs apply")
             var = await self.new_entity(conf, *args, **kwargs)
             if self.parent is not None:
                 await cg.register_parented(var, self.parent)
         else:
-            fn = await self._apply_function(conf[CONF_ID], apply, publish)
+            fn = await self._apply_function(entity_id, apply, publish)
             var = await self.new_entity(conf, cg.TemplateArguments(fn), *args, **kwargs)
         if setter is not None:
             add(setter(var))
         return var
 
     async def _apply_function(
-        self, entity_id: ID, apply: str, publish: str | None
+        self, entity_id: ID, apply: str, publish: ApplyPublish
     ) -> SafeExpType:
-        if self.apply_entity is None or self.parent is None:
-            raise ValueError(
-                f"{entity_id}: apply needs a parent and a domain that supports it"
-            )
-        if publish is not None and (
-            publish not in ("before", "after") or self.apply_entity.publish is None
-        ):
-            raise ValueError(f"{entity_id}: cannot publish {publish!r} for this domain")
         parent = self.parent
         if isinstance(parent, ID):
             parent = await get_variable(parent)
-        # Named from global scope, so a parameter cannot shadow the hub.
-        body = [f"::{parent}->{apply};"]
+        body = [f"{parent_ref(parent)}->{apply};"]
         if publish == "before":
             body.insert(0, self.apply_entity.publish)
         elif publish == "after":

tests/unit_tests/core/test_entity_helpers.py

@@ -1,13 +1,14 @@
 """Test get_base_entity_object_id function matches C++ behavior."""
 
 from collections.abc import Callable, Generator
-import importlib
 from pathlib import Path
 import re
 from typing import Any
 
 import pytest
 
+import esphome.codegen as cg
+from esphome.components import button, number, select, switch
 from esphome.config_validation import Invalid
 from esphome.const import (
     CONF_DEVICE_CLASS,
@@ -34,7 +35,7 @@ from esphome.core.entity_helpers import (
     setup_entity,
     setup_unit_of_measurement,
 )
-from esphome.cpp_generator import MockObj, RawExpression
+from esphome.cpp_generator import MockObj, MockObjClass, RawExpression
 from esphome.helpers import fnv1_hash, sanitize, snake_case
 
 from .common import load_config_from_fixture
@@ -1256,10 +1257,11 @@ async def test_sub_entities_creates_and_sets(
 
     monkeypatch.setattr(entity_helpers.cg, "register_parented", fake_register_parented)
     hub = MockObj("hub", "->")
-    sub = entity_helpers.SubEntities(new_entity, {"child": {}}, parent=hub)
+    conf = {CONF_ID: ID("child")}
+    sub = entity_helpers.SubEntities(new_entity, {"child": conf}, parent=hub)
 
     assert await sub("child", hub.set_child, 1, step=2) is created
-    assert calls == [({}, (1,), {"step": 2})]
+    assert calls == [(conf, (1,), {"step": 2})]
     assert setup_test_environment == ["parent child hub", "hub->set_child(child)"]
 
 
@@ -1315,24 +1317,33 @@ def apply_functions(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, str, str
     return emitted
 
 
+_HUB = MockObj("hub", "->")
+_APPLY_SWITCH_CLS = cg.esphome_ns.namespace("switch_").class_("ApplySwitch")
 _SWITCH_APPLY = entity_helpers.ApplyEntity(
+    _APPLY_SWITCH_CLS,
     ((RawExpression("switch_::Switch *"), "entity"), (RawExpression("bool"), "state")),
     "entity->publish_state(state);",
 )
+_BUTTON_APPLY = entity_helpers.ApplyEntity(
+    cg.esphome_ns.namespace("button").class_("ApplyButton"),
+    ((RawExpression("button::Button *"), "entity"),),
+)
+
+
+def _apply_conf(name: str, type_: MockObjClass = _APPLY_SWITCH_CLS) -> dict[str, ID]:
+    return {CONF_ID: ID(name, type=type_)}
 
 
 def _apply_sub(
     calls: list[tuple[Any, ...]],
     config: dict[str, Any],
-    apply_entity: entity_helpers.ApplyEntity | None = _SWITCH_APPLY,
-    parent: MockObj | ID | None = None,
+    apply_entity: entity_helpers.ApplyEntity = _SWITCH_APPLY,
+    parent: MockObj | ID | None = _HUB,
 ) -> entity_helpers.SubEntities:
     async def new_entity(conf: dict[str, Any], *args: Any, **kwargs: Any) -> MockObj:
         calls.append((conf[CONF_ID].id, *(str(a) for a in args)))
         return MockObj(conf[CONF_ID].id)
 
-    if parent is None:
-        parent = MockObj("hub", "->")
     return entity_helpers.SubEntities(new_entity, config, parent, apply_entity)
 
 
@@ -1363,7 +1374,7 @@ async def test_sub_entities_apply_generates_function(
     monkeypatch.setattr(entity_helpers.cg, "register_parented", fail_register_parented)
     calls: list[tuple[Any, ...]] = []
     hub = MockObj("hub", "->")
-    sub = _apply_sub(calls, {"bt": {CONF_ID: ID("bt_sw")}})
+    sub = _apply_sub(calls, {"bt": _apply_conf("bt_sw")})
 
     var = await sub(
         "bt", hub.set_bt_switch, apply="set_bluetooth(state)", publish=publish
@@ -1384,9 +1395,9 @@ async def test_sub_entities_apply_shares_equal_functions(
     calls: list[tuple[Any, ...]] = []
     hub = MockObj("hub", "->")
     config = {
-        "a": {CONF_ID: ID("a_sw")},
-        "b": {CONF_ID: ID("b_sw")},
-        "c": {CONF_ID: ID("c_sw")},
+        "a": _apply_conf("a_sw"),
+        "b": _apply_conf("b_sw"),
+        "c": _apply_conf("c_sw"),
     }
     sub = _apply_sub(calls, config)
     await sub("a", hub.set_a, apply="refresh()", publish="before")
@@ -1417,66 +1428,98 @@ async def test_sub_entities_apply_parent_id(
 
     monkeypatch.setattr(entity_helpers, "get_variable", fake_get_variable)
     calls: list[tuple[Any, ...]] = []
-    sub = _apply_sub(calls, {"r": {CONF_ID: ID("r_btn")}}, parent=ID("my_hub"))
+    sub = _apply_sub(calls, {"r": _apply_conf("r_sw")}, parent=ID("my_hub"))
     await sub("r", MockObj("my_hub", "->").set_r, apply="restart()")
     assert apply_functions[0][2] == "::my_hub->restart();"
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("apply_entity", "kwargs", "match"),
+    ("apply_entity", "conf", "kwargs", "match"),
     [
-        (None, {"apply": "restart()"}, "apply needs"),
+        (_SWITCH_APPLY, _apply_conf("k_id"), {}, "needs apply="),
         (
-            entity_helpers.ApplyEntity(((RawExpression("button::Button *"), "e"),)),
+            _SWITCH_APPLY,
+            _apply_conf("k_id", cg.esphome_ns.namespace("x").class_("MySwitch")),
+            {"apply": "x()"},
+            "needs the domain's Apply class",
+        ),
+        (
+            _BUTTON_APPLY,
+            _apply_conf(
+                "k_id", cg.esphome_ns.namespace("button").class_("ApplyButton")
+            ),
             {"apply": "restart()", "publish": "before"},
-            "cannot publish",
+            "publish= needs",
+        ),
+        (
+            _SWITCH_APPLY,
+            _apply_conf("k_id", cg.esphome_ns.namespace("x").class_("MySwitch")),
+            {"publish": "before"},
+            "publish= needs",
         ),
-        (_SWITCH_APPLY, {"apply": "x()", "publish": "sometimes"}, "cannot publish"),
-        (_SWITCH_APPLY, {"publish": "before"}, "publish needs apply"),
     ],
 )
 async def test_sub_entities_apply_errors(
     setup_test_environment: list[str],
     apply_functions: list[tuple[str, str, str]],
-    apply_entity: entity_helpers.ApplyEntity | None,
+    apply_entity: entity_helpers.ApplyEntity,
+    conf: dict[str, ID],
     kwargs: dict[str, str],
     match: str,
 ) -> None:
+    """A mismatch between the declared class and apply= fails in Python, not in C++."""
     calls: list[tuple[Any, ...]] = []
-    sub = _apply_sub(calls, {"k": {CONF_ID: ID("k_id")}}, apply_entity)
+    sub = _apply_sub(calls, {"k": conf}, apply_entity)
     with pytest.raises(ValueError, match=match):
         await sub("k", MockObj("hub", "->").set_k, **kwargs)
     assert calls == []
     assert apply_functions == []
 
 
+@pytest.mark.asyncio
+async def test_sub_entities_apply_needs_parent(
+    setup_test_environment: list[str], apply_functions: list[tuple[str, str, str]]
+) -> None:
+    calls: list[tuple[Any, ...]] = []
+    sub = _apply_sub(calls, {"k": _apply_conf("k_id")}, parent=None)
+    with pytest.raises(ValueError, match="needs a parent"):
+        await sub("k", None, apply="x()")
+
+
 @pytest.mark.parametrize(
-    ("sub_fn", "params", "publish"),
+    ("sub_fn", "cls", "params", "publish"),
     [
-        ("button.sub_buttons", "button::Button * entity", None),
+        (button.sub_buttons, button.ApplyButton, "button::Button * entity", None),
         (
-            "switch.sub_switches",
+            switch.sub_switches,
+            switch.ApplySwitch,
             "switch_::Switch * entity, bool state",
             "entity->publish_state(state);",
         ),
         (
-            "number.sub_numbers",
+            number.sub_numbers,
+            number.ApplyNumber,
             "number::Number * entity, float value",
             "entity->publish_state(value);",
         ),
         (
-            "select.sub_selects",
+            select.sub_selects,
+            select.ApplySelect,
             "select::Select * entity, size_t index",
             "entity->publish_state(index);",
         ),
     ],
 )
-def test_domain_apply_entities(sub_fn: str, params: str, publish: str | None) -> None:
-    """Each domain's sub helper declares the generated function its Apply class calls."""
-    module, name = sub_fn.split(".")
-    imported = importlib.import_module(f"esphome.components.{module}")
-    apply_entity = getattr(imported, name)({}).apply_entity
+def test_domain_apply_entities(
+    sub_fn: Callable[..., entity_helpers.SubEntities],
+    cls: MockObjClass,
+    params: str,
+    publish: str | None,
+) -> None:
+    """Each domain's sub helper declares its Apply class and the function it calls."""
+    apply_entity = sub_fn({}).apply_entity
+    assert apply_entity.cls is cls
     assert ", ".join(f"{type_} {arg}" for type_, arg in apply_entity.params) == params
     assert apply_entity.publish == publish
 
@@ -1487,9 +1530,9 @@ async def test_sub_entities_without_setter(
 ) -> None:
     """A hub that does not keep the entity passes no setter; nothing else is emitted."""
     calls: list[tuple[Any, ...]] = []
-    sub = _apply_sub(calls, {"vent": {CONF_ID: ID("vent_btn")}})
+    sub = _apply_sub(calls, {"vent": _apply_conf("vent_sw")})
     assert await sub("vent", None, apply="vent_door()") is not None
-    assert calls == [("vent_btn", "<esphome__vent_btn__fn>")]
+    assert calls == [("vent_sw", "<esphome__vent_sw__fn>")]
     assert setup_test_environment == []
 
 
@@ -1499,11 +1542,13 @@ async def test_sub_entities_create_from_whole_config(
 ) -> None:
     """create() builds the entity from a config that is the entity itself."""
     calls: list[tuple[Any, ...]] = []
-    conf = {CONF_ID: ID("rf_sw")}
     sub = _apply_sub(calls, {})
     hub = MockObj("hub", "->")
     var = await sub.create(
-        conf, hub.set_rf_switch, apply="set_rf_mode(state)", publish="before"
+        _apply_conf("rf_sw"),
+        hub.set_rf_switch,
+        apply="set_rf_mode(state)",
+        publish="before",
     )
     assert str(var) == "rf_sw"
     assert calls == [("rf_sw", "<esphome__rf_sw__fn>")]