[display] Add add_writer for display lambdas (#21027)

esphome/components/display/__init__.py

@@ -27,6 +27,7 @@ from esphome.const import (
 )
 from esphome.core import CORE, ID, CoroPriority, coroutine_with_priority
 from esphome.final_validate import full_config
+from esphome.types import ConfigType
 
 DOMAIN = "display"
 IS_PLATFORM_COMPONENT = True
@@ -291,6 +292,20 @@ async def register_display(var, config):
     await setup_display_core_(var, config)
 
 
+async def add_writer(
+    var: cg.MockObj, config: ConfigType, ref: cg.MockObj = DisplayRef
+) -> None:
+    """Pass the configured lambda, if any, to var.set_writer.
+
+    ref is the type of `it`; drivers with their own drawing class pass theirs.
+    """
+    if (lambda_config := config.get(CONF_LAMBDA)) is not None:
+        lambda_ = await cg.process_lambda(
+            lambda_config, [(ref, "it")], return_type=cg.void
+        )
+        cg.add(var.set_writer(lambda_))
+
+
 @automation.register_action(
     "display.page.show",
     DisplayPageShowAction,

tests/unit_tests/components/test_display.py

@@ -0,0 +1,41 @@
+"""Tests for display.add_writer."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from esphome import codegen as cg
+from esphome.components import display
+from esphome.const import CONF_LAMBDA
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "ref"),
+    [({}, display.DisplayRef), ({"ref": "MyDisplayRef"}, "MyDisplayRef")],
+)
+@pytest.mark.asyncio
+async def test_add_writer_sets_the_lambda(kwargs: dict, ref: object) -> None:
+    var = MagicMock(name="var")
+    lambda_config = MagicMock(name="lambda")
+    with (
+        patch.object(cg, "process_lambda", AsyncMock(return_value="fn")) as process,
+        patch.object(cg, "add") as add,
+    ):
+        await display.add_writer(var, {CONF_LAMBDA: lambda_config}, **kwargs)
+
+    process.assert_awaited_once_with(lambda_config, [(ref, "it")], return_type=cg.void)
+    var.set_writer.assert_called_once_with("fn")
+    add.assert_called_once_with(var.set_writer.return_value)
+
+
+@pytest.mark.asyncio
+async def test_add_writer_without_lambda_does_nothing() -> None:
+    var = MagicMock(name="var")
+    with (
+        patch.object(cg, "process_lambda", AsyncMock()) as process,
+        patch.object(cg, "add") as add,
+    ):
+        await display.add_writer(var, {})
+
+    process.assert_not_awaited()
+    add.assert_not_called()