[core] Add ApplyBytes for apply actions that forward a byte payload

esphome/automation.py

@@ -356,6 +356,27 @@ class ApplyField:
         return ApplyCall(target, ((self.conf_key, self.type_, self.const_fn),))
 
 
+@dataclass(frozen=True)
+class ApplyBytes:
+    """A byte payload key forwarded as ``target``, whose ``{}`` stands for ``data, len``.
+
+    A constant payload is a shared PROGMEM table passed as is; on ESP8266 it is first copied
+    to the stack (``stack`` bytes, heap beyond that) unless ``progmem_target`` names a call
+    that reads flash itself. A lambda's ``std::vector<uint8_t>`` is passed as its data and
+    size. The template may name the parent as ``{parent}``. An absent key emits nothing.
+    """
+
+    conf_key: str | tuple[str, ...]
+    target: str
+    stack: int = 64
+    progmem_target: str | None = None
+
+    def __post_init__(self) -> None:
+        for target in (self.target, self.progmem_target):
+            if target is not None and target.count("{}") != 1:
+                raise ValueError(f"apply bytes target {target!r} needs one {{}}")
+
+
 def _config_lookup(config: ConfigType, key: str | tuple[str, ...]) -> Any:
     if isinstance(key, str):
         return config.get(key)
@@ -488,7 +509,7 @@ def _apply_values(config: ConfigType, members: list[tuple[Any, Any, Any]]) -> li
 def register_apply_action(
     name: str,
     schema: cv.Schema,
-    *fields: ApplyField | ApplyCall,
+    *fields: ApplyField | ApplyCall | ApplyBytes,
     call: str | None = None,
     id_key: str = CONF_ID,
 ) -> None:
@@ -504,7 +525,7 @@ def register_apply_action(
 def apply_action_builder(
     name: str,
     schema: cv.Schema,
-    *fields: ApplyField | ApplyCall,
+    *fields: ApplyField | ApplyCall | ApplyBytes,
     call: str | None = None,
     id_key: str = CONF_ID,
 ) -> Callable[..., Awaitable[MockObj]]:
@@ -517,16 +538,20 @@ def apply_action_builder(
     and ``apply_call.perform()`` is appended.
     """
     # An action stores the value, so a std::string constant stays in flash on ESP8266.
-    statements_spec = [
-        (
-            c.target,
-            _default_std_string(c.members, flash_string),
-        )
-        for c in (f if isinstance(f, ApplyCall) else f.call() for f in fields)
-    ]
+    statements_spec: list[ApplyBytes | tuple[str, list[tuple[Any, Any, Any]]]] = []
+    for f in fields:
+        if isinstance(f, ApplyBytes):
+            statements_spec.append(f)
+            continue
+        c = f if isinstance(f, ApplyCall) else f.call()
+        statements_spec.append((c.target, _default_std_string(c.members, flash_string)))
+    table_name = name.replace(".", "_")
     _check_key_in_schema(name, schema, id_key)
-    for _, members in statements_spec:
-        for conf_key, _, _ in members:
+    for spec in statements_spec:
+        if isinstance(spec, ApplyBytes):
+            _check_key_in_schema(name, schema, spec.conf_key)
+            continue
+        for conf_key, _, _ in spec[1]:
             _check_key_in_schema(name, schema, conf_key)
 
     async def builder(
@@ -539,7 +564,15 @@ def apply_action_builder(
         lambda_args = _apply_lambda_args(args)
         receiver = "apply_call." if call else f"{parent}->"
         statements: list[str] = []
-        for target, members in statements_spec:
+        for spec in statements_spec:
+            if isinstance(spec, ApplyBytes):
+                statements.extend(
+                    await _render_bytes(
+                        spec, config, receiver, parent, lambda_args, table_name
+                    )
+                )
+                continue
+            target, members = spec
             values = _apply_values(config, members)
             if not _apply_call_active(members, values):
                 continue
@@ -560,6 +593,49 @@ def apply_action_builder(
     return builder
 
 
+async def _render_bytes(
+    spec: ApplyBytes,
+    config: ConfigType,
+    receiver: str,
+    parent: str,
+    lambda_args: TemplateArgsType,
+    table_name: str,
+) -> list[str]:
+    """The statements for one ``ApplyBytes``: what ``TemplatableBytes::visit`` does, inlined."""
+    value = _config_lookup(config, spec.conf_key)
+    if value is None:
+        return []
+
+    def send(target: str, data: str) -> str:
+        return f"{receiver}{target.format(data, parent=parent)};"
+
+    if isinstance(value, Lambda):
+        inner = await cg.process_lambda(
+            value, lambda_args, return_type=cg.std_vector.template(cg.uint8)
+        )
+        return [
+            "{",
+            f"  const std::vector<uint8_t> apply_bytes = {call_lambda(inner)};",
+            f"  {send(spec.target, 'apply_bytes.data(), apply_bytes.size()')}",
+            "}",
+        ]
+    if (size := len(value)) > 0xFFFF:
+        raise EsphomeError(f"Byte payload is {size} bytes; the maximum is 65535")
+    table = f"{progmem_bytes(table_name, value)}, {size}"
+    if not CORE.is_esp8266 or size == 0:
+        return [send(spec.target, table)]
+    if spec.progmem_target is not None:
+        return [send(spec.progmem_target, table)]
+    # A PROGMEM table must be read with aligned loads on ESP8266, so copy it out first
+    return [
+        "{",
+        f"  SmallBufferWithHeapFallback<{spec.stack}> apply_bytes({size});",
+        f"  progmem_memcpy(apply_bytes.get(), {table});",
+        f"  {send(spec.target, f'apply_bytes.get(), {size}')}",
+        "}",
+    ]
+
+
 def _apply_call_active(members: list[tuple[Any, Any, Any]], values: list[Any]) -> bool:
     """An ``ApplyCall`` is emitted unless it has keys and none of them is set."""
     return not members or any(value is not None for value in values)

tests/unit_tests/test_automation.py

@@ -11,6 +11,7 @@ from esphome.automation import (
     CONF_CHOICE_ID,
     CONF_CHOOSE,
     ApplyAction,
+    ApplyBytes,
     ApplyCall,
     ApplyCondition,
     ApplyField,
@@ -1281,6 +1282,116 @@ async def test_apply_string_ref_literal_stays_in_flash_on_esp8266(
     assert f"::{PARENT_OBJ}->set_payload({rendered});" in _apply_definition(mock_cg)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("platform", ["esp32", "esp8266"])
+async def test_apply_bytes_lambda_is_a_vector_in_a_block(
+    registries: tuple[Registry, Registry], mock_cg: MockCodegen, platform: str
+) -> None:
+    """A lambda payload is called once into a scoped vector and passed as data and size."""
+    fields = (ApplyBytes("data", "send({})"),)
+    config = {"data": Lambda("return {0x01, x};")}
+    await _run_apply_action(
+        registries, fields, config, args=[(cg.uint8, "x")], platform=platform
+    )
+    text = _apply_definition(mock_cg)
+    assert "const std::vector<uint8_t> apply_bytes = " in text
+    assert "-> std::vector<uint8_t>" in text
+    assert f"::{PARENT_OBJ}->send(apply_bytes.data(), apply_bytes.size());" in text
+    assert "progmem_memcpy" not in text
+
+
+@pytest.mark.asyncio
+async def test_apply_bytes_constant_is_a_shared_table(
+    registries: tuple[Registry, Registry], mock_cg: MockCodegen
+) -> None:
+    """A constant payload is the shared PROGMEM table and its length, named after the action."""
+    CORE.config = {}
+    fields = (ApplyBytes("data", "send({})"),)
+    await _run_apply_action(registries, fields, {"data": [0xA1, 0x02]})
+    assert f"::{PARENT_OBJ}->send(my_apply, 2);" in _apply_definition(mock_cg)
+
+    mock_cg.new_pvariable.reset_mock()
+    await _run_apply_action(registries, fields, {"data": b""})
+    assert f"::{PARENT_OBJ}->send(nullptr, 0);" in _apply_definition(mock_cg)
+
+    mock_cg.new_pvariable.reset_mock()
+    await _run_apply_action(registries, fields, {})
+    assert "send(" not in _apply_definition(mock_cg)
+
+    with pytest.raises(EsphomeError, match="maximum is 65535"):
+        await _run_apply_action(registries, fields, {"data": bytes(0x10000)})
+
+
+@pytest.mark.asyncio
+async def test_apply_bytes_constant_is_copied_out_of_flash_on_esp8266(
+    registries: tuple[Registry, Registry], mock_cg: MockCodegen
+) -> None:
+    """On ESP8266 a constant goes through a stack buffer, or straight to a flash-aware target."""
+    CORE.config = {}
+    fields = (ApplyBytes("data", "send({})", stack=128),)
+    await _run_apply_action(
+        registries, fields, {"data": [0xA1, 0x02]}, platform="esp8266"
+    )
+    text = _apply_definition(mock_cg)
+    lines = [
+        "SmallBufferWithHeapFallback<128> apply_bytes(2);",
+        "progmem_memcpy(apply_bytes.get(), my_apply, 2);",
+        f"::{PARENT_OBJ}->send(apply_bytes.get(), 2);",
+    ]
+    positions = [text.index(line) for line in lines]
+    assert positions == sorted(positions)
+
+    mock_cg.new_pvariable.reset_mock()
+    fields = (ApplyBytes("data", "send({})", progmem_target="send_progmem({})"),)
+    await _run_apply_action(
+        registries, fields, {"data": [0xA1, 0x02]}, platform="esp8266"
+    )
+    text = _apply_definition(mock_cg)
+    assert f"::{PARENT_OBJ}->send_progmem(my_apply, 2);" in text
+    assert "progmem_memcpy" not in text
+
+    # An empty payload never touches flash
+    mock_cg.new_pvariable.reset_mock()
+    await _run_apply_action(registries, fields, {"data": []}, platform="esp8266")
+    assert f"::{PARENT_OBJ}->send(nullptr, 0);" in _apply_definition(mock_cg)
+
+
+@pytest.mark.asyncio
+async def test_apply_bytes_beside_other_fields_and_call_object(
+    registries: tuple[Registry, Registry], mock_cg: MockCodegen
+) -> None:
+    CORE.config = {}
+    fields = (
+        ApplyField("channel", "set_channel", cg.uint8),
+        ApplyBytes("data", "write({})"),
+    )
+    await _run_apply_action(
+        registries, fields, {"channel": 3, "data": [1]}, call="make_call"
+    )
+    text = _apply_definition(mock_cg)
+    lines = [
+        f"auto apply_call = ::{PARENT_OBJ}->make_call();",
+        "apply_call.set_channel(3);",
+        "apply_call.write(my_apply, 1);",
+        "apply_call.perform();",
+    ]
+    positions = [text.index(line) for line in lines]
+    assert positions == sorted(positions)
+
+
+def test_apply_bytes_checks_its_shape(registries: tuple[Registry, Registry]) -> None:
+    with pytest.raises(ValueError, match="needs one {}"):
+        ApplyBytes("data", "send()")
+    with pytest.raises(ValueError, match="needs one {}"):
+        ApplyBytes("data", "send({})", progmem_target="send_progmem({}, {})")
+    with pytest.raises(ValueError, match="not in the schema"):
+        register_apply_action(
+            "my.bad",
+            cv.Schema({cv.Required("data"): cv.string}),
+            ApplyBytes("dat", "send({})"),
+        )
+
+
 @pytest.mark.asyncio
 async def test_register_apply_condition_predicate(
     registries: tuple[Registry, Registry], mock_cg: MockCodegen