[core] Simplify the Apply entity helpers

Fold the function cache into SubEntities, make each domain's APPLY_ENTITY public,
shorten the class and docstring comments and reuse the shared test hub.

esphome/components/button/__init__.py

@@ -118,14 +118,14 @@ async def new_button(config: ConfigType, *args: SafeExpType) -> MockObj:
     return var
 
 
-_APPLY_ENTITY = ApplyEntity(ApplyButton, ((ButtonPtr, "entity"),))
+APPLY_ENTITY = ApplyEntity(ApplyButton, ((ButtonPtr, "entity"),))
 
 
 def sub_buttons(
     config: ConfigType, *, parent: MockObj | ID | None = None
 ) -> SubEntities:
     """Return a SubEntities bound to new_button; ``apply`` statements can use ``entity``."""
-    return SubEntities(new_button, config, parent, _APPLY_ENTITY)
+    return SubEntities(new_button, config, parent, APPLY_ENTITY)
 
 
 BUTTON_PRESS_SCHEMA = maybe_simple_id(

esphome/components/button/button.h

@@ -46,11 +46,10 @@ class Button : public EntityBase {
   LazyCallbackManager<void()> press_callback_{};
 };
 
-/// Runs one codegen-generated function on a press, with the hub and any publish baked in.
-/// Fn takes the entity first, so the function can publish or read options without a stored pointer.
+/// Calls the codegen-generated Fn(this) on a press; the hub and any publish are baked into Fn.
 template<auto Fn> class ApplyButton final : public Button {
  public:
-  // User provided, not "= default": `new(p) ApplyButton()` would zero-fill .bss that is already zero.
+  // Not "= default": placement new would zero-fill .bss that is already zero.
   ApplyButton() {}
 
  protected:

esphome/components/number/__init__.py

@@ -330,7 +330,7 @@ async def new_number(config, *args, min_value: float, max_value: float, step: fl
     return var
 
 
-_APPLY_ENTITY = ApplyEntity(
+APPLY_ENTITY = ApplyEntity(
     ApplyNumber,
     ((NumberPtr, "entity"), (cg.float_, "value")),
     "entity->publish_state(value);",
@@ -341,7 +341,7 @@ def sub_numbers(
     config: ConfigType, *, parent: MockObj | ID | None = None
 ) -> SubEntities:
     """Return a SubEntities bound to new_number; ``apply`` statements can use ``entity`` and ``value``."""
-    return SubEntities(new_number, config, parent, _APPLY_ENTITY)
+    return SubEntities(new_number, config, parent, APPLY_ENTITY)
 
 
 NUMBER_IN_RANGE_CONDITION_SCHEMA = cv.All(

esphome/components/number/number.h

@@ -54,11 +54,10 @@ class Number : public EntityBase {
   LazyCallbackManager<void(float)> state_callback_;
 };
 
-/// Runs one codegen-generated function on a new value, with the hub and any publish baked in.
-/// Fn takes the entity first, so the function can publish or read options without a stored pointer.
+/// Calls the codegen-generated Fn(this, value) on a new value; the hub and any publish are baked into Fn.
 template<auto Fn> class ApplyNumber final : public Number {
  public:
-  // User provided, not "= default": `new(p) ApplyNumber()` would zero-fill .bss that is already zero.
+  // Not "= default": placement new would zero-fill .bss that is already zero.
   ApplyNumber() {}
 
  protected:

esphome/components/select/__init__.py

@@ -126,7 +126,7 @@ async def new_select(config, *args, options: list[str]):
     return var
 
 
-_APPLY_ENTITY = ApplyEntity(
+APPLY_ENTITY = ApplyEntity(
     ApplySelect,
     ((SelectPtr, "entity"), (cg.size_t, "index")),
     "entity->publish_state(index);",
@@ -137,7 +137,7 @@ 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 SubEntities(new_select, config, parent, _APPLY_ENTITY)
+    return SubEntities(new_select, config, parent, APPLY_ENTITY)
 
 
 @coroutine_with_priority(CoroPriority.CORE)

esphome/components/select/select.h

@@ -112,11 +112,10 @@ class Select : public EntityBase {
   LazyCallbackManager<void(size_t)> state_callback_;
 };
 
-/// Runs one codegen-generated function on a new index, with the hub and any publish baked in.
-/// Fn takes the entity first, so the function can publish or read options without a stored pointer.
+/// Calls the codegen-generated Fn(this, index) on a new index; the hub and any publish are baked into Fn.
 template<auto Fn> class ApplySelect final : public Select {
  public:
-  // User provided, not "= default": `new(p) ApplySelect()` would zero-fill .bss that is already zero.
+  // Not "= default": placement new would zero-fill .bss that is already zero.
   ApplySelect() {}
 
  protected:

esphome/components/switch/__init__.py

@@ -178,7 +178,7 @@ async def new_switch(config, *args):
     return var
 
 
-_APPLY_ENTITY = ApplyEntity(
+APPLY_ENTITY = ApplyEntity(
     ApplySwitch,
     ((SwitchPtr, "entity"), (cg.bool_, "state")),
     "entity->publish_state(state);",
@@ -189,7 +189,7 @@ def sub_switches(
     config: ConfigType, *, parent: MockObj | ID | None = None
 ) -> SubEntities:
     """Return a SubEntities bound to new_switch; ``apply`` statements can use ``entity`` and ``state``."""
-    return SubEntities(new_switch, config, parent, _APPLY_ENTITY)
+    return SubEntities(new_switch, config, parent, APPLY_ENTITY)
 
 
 SWITCH_ACTION_SCHEMA = maybe_simple_id(

esphome/components/switch/switch.h

@@ -147,11 +147,10 @@ class Switch : public EntityBase {
 #define LOG_SWITCH(prefix, type, obj) log_switch((TAG), (prefix), LOG_STR_LITERAL(type), (obj))
 void log_switch(const char *tag, const char *prefix, const char *type, Switch *obj);
 
-/// Runs one codegen-generated function on a state write, with the hub and any publish baked in.
-/// Fn takes the entity first, so the function can publish or read options without a stored pointer.
+/// Calls the codegen-generated Fn(this, state) on a state write; the hub and any publish are baked into Fn.
 template<auto Fn> class ApplySwitch final : public Switch {
  public:
-  // User provided, not "= default": `new(p) ApplySwitch()` would zero-fill .bss that is already zero.
+  // Not "= default": placement new would zero-fill .bss that is already zero.
   ApplySwitch() {}
 
  protected:

esphome/core/entity_helpers.py

@@ -484,39 +484,11 @@ class ApplyEntity:
 _APPLY_DOMAIN = "entity_apply"
 
 
-@dataclass
-class _ApplyData:
-    apply_functions: dict[str, SafeExpType] = field(default_factory=dict)
-
-
-def _get_apply_data() -> _ApplyData:
-    if _APPLY_DOMAIN not in CORE.data:
-        CORE.data[_APPLY_DOMAIN] = _ApplyData()
-    return CORE.data[_APPLY_DOMAIN]
-
-
-def _apply_entity_function(
-    entity_id: ID, params: tuple[tuple[SafeExpType, str], ...], body: list[str]
-) -> SafeExpType:
-    """Emit the control function once per distinct body, so equal entities share one class."""
-    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] = apply_static_function(
-            entity_id, cg.void, list(params), body
-        )
-    return fn
-
-
 @dataclass(frozen=True, slots=True)
 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.
+    Returns None only when key is absent. ``apply``/``publish`` are documented on ``create``.
     """
 
     new_entity: Callable[..., Awaitable[MockObj]]
@@ -528,7 +500,7 @@ class SubEntities:
         self,
         key: str,
         setter: Callable[[MockObj], Expression] | None,
-        *args: Any,
+        *args: SafeExpType,
         **kwargs: Any,
     ) -> MockObj | None:
         if (conf := self.config.get(key)) is None:
@@ -544,7 +516,11 @@ class SubEntities:
         publish: ApplyPublish = None,
         **kwargs: Any,
     ) -> MockObj:
-        """Create the entity from ``conf`` itself, for a platform whose config is one entity."""
+        """Create the entity from ``conf``; ``setter`` may be None when the hub keeps no pointer.
+
+        With ``apply`` the entity is the domain's ``Apply*`` class and its control runs
+        ``parent->apply``, with ``publish`` ("before"/"after") placing the domain's publish.
+        """
         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
@@ -576,15 +552,23 @@ class SubEntities:
     async def _apply_function(
         self, entity_id: ID, apply: str, publish: ApplyPublish
     ) -> SafeExpType:
+        """Emit the control function once per distinct body, so equal entities share one class."""
         parent = self.parent
         if isinstance(parent, ID):
             parent = await get_variable(parent)
+        apply_entity = self.apply_entity
         body = [f"{parent_ref(parent)}->{apply};"]
         if publish == "before":
-            body.insert(0, self.apply_entity.publish)
+            body.insert(0, apply_entity.publish)
         elif publish == "after":
-            body.append(self.apply_entity.publish)
-        return _apply_entity_function(entity_id, self.apply_entity.params, body)
+            body.append(apply_entity.publish)
+        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:
+            fn = functions[key] = apply_static_function(
+                entity_id, cg.void, list(apply_entity.params), body
+            )
+        return fn
 
 
 def inherit_property_from(property_to_inherit, parent_id_property, transform=None):

tests/unit_tests/core/test_entity_helpers.py

@@ -1373,11 +1373,10 @@ 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": _apply_conf("bt_sw")})
 
     var = await sub(
-        "bt", hub.set_bt_switch, apply="set_bluetooth(state)", publish=publish
+        "bt", _HUB.set_bt_switch, apply="set_bluetooth(state)", publish=publish
     )
     assert str(var) == "bt_sw"
     assert calls == [("bt_sw", "<esphome__bt_sw__fn>")]
@@ -1393,16 +1392,15 @@ async def test_sub_entities_apply_shares_equal_functions(
 ) -> None:
     """Entities with the same statement share one function, so one class instantiation."""
     calls: list[tuple[Any, ...]] = []
-    hub = MockObj("hub", "->")
     config = {
         "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")
-    await sub("b", hub.set_b, apply="refresh()", publish="before")
-    await sub("c", hub.set_c, apply="refresh()")
+    await sub("a", _HUB.set_a, apply="refresh()", publish="before")
+    await sub("b", _HUB.set_b, apply="refresh()", publish="before")
+    await sub("c", _HUB.set_c, apply="refresh()")
     assert calls == [
         ("a_sw", "<esphome__a_sw__fn>"),
         ("b_sw", "<esphome__a_sw__fn>"),
@@ -1549,10 +1547,9 @@ async def test_sub_entities_create_from_whole_config(
     """create() builds the entity from a config that is the entity itself."""
     calls: list[tuple[Any, ...]] = []
     sub = _apply_sub(calls, {})
-    hub = MockObj("hub", "->")
     var = await sub.create(
         _apply_conf("rf_sw"),
-        hub.set_rf_switch,
+        _HUB.set_rf_switch,
         apply="set_rf_mode(state)",
         publish="before",
     )