From 2fe1333cba6d8efc8219153b992f238076d2bc4b Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 17 Oct 2026 18:34:38 -1000
Subject: [PATCH] [display] Add add_writer for display lambdas (#21027)

---
 esphome/components/display/__init__.py      | 15 ++++++++
 tests/unit_tests/components/test_display.py | 41 +++++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 tests/unit_tests/components/test_display.py

diff --git a/esphome/components/display/__init__.py b/esphome/components/display/__init__.py
index 22e3d9320b0..be8586b3b25 100644
--- a/esphome/components/display/__init__.py
+++ b/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,
diff --git a/tests/unit_tests/components/test_display.py b/tests/unit_tests/components/test_display.py
new file mode 100644
index 00000000000..a98dd80ec0e
--- /dev/null
+++ b/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()