From c42343be3ae8082b804d9633446441192bde3fc6 Mon Sep 17 00:00:00 2001
From: Clyde Stubbs <2366188+clydebarrow@users.noreply.github.com>
Date: Sun, 30 Mar 2025 20:09:19 +1100
Subject: [PATCH] [lvgl] Implement switch restore (#8481)

---
 esphome/components/lvgl/switch/__init__.py   | 21 ++++++++++++--------
 esphome/components/lvgl/switch/lvgl_switch.h | 21 +++++---------------
 2 files changed, 18 insertions(+), 24 deletions(-)

diff --git a/esphome/components/lvgl/switch/__init__.py b/esphome/components/lvgl/switch/__init__.py
index 4e1e7f72e0..6d10a70d85 100644
--- a/esphome/components/lvgl/switch/__init__.py
+++ b/esphome/components/lvgl/switch/__init__.py
@@ -1,7 +1,9 @@
 import esphome.codegen as cg
-from esphome.components.switch import Switch, new_switch, switch_schema
+from esphome.components.switch import Switch, register_switch, switch_schema
 import esphome.config_validation as cv
+from esphome.const import CONF_ID
 from esphome.cpp_generator import MockObj
+from esphome.cpp_types import Component
 
 from ..defines import CONF_WIDGET, literal
 from ..lvcode import (
@@ -18,7 +20,7 @@ from ..lvcode import (
 from ..types import LV_EVENT, LV_STATE, lv_pseudo_button_t, lvgl_ns
 from ..widgets import get_widgets, wait_for_widgets
 
-LVGLSwitch = lvgl_ns.class_("LVGLSwitch", Switch)
+LVGLSwitch = lvgl_ns.class_("LVGLSwitch", Switch, Component)
 CONFIG_SCHEMA = switch_schema(LVGLSwitch).extend(
     {
         cv.Required(CONF_WIDGET): cv.use_id(lv_pseudo_button_t),
@@ -27,21 +29,24 @@ CONFIG_SCHEMA = switch_schema(LVGLSwitch).extend(
 
 
 async def to_code(config):
-    switch = await new_switch(config)
     widget = await get_widgets(config, CONF_WIDGET)
     widget = widget[0]
     await wait_for_widgets()
-    async with LambdaContext(EVENT_ARG) as checked_ctx:
-        checked_ctx.add(switch.publish_state(widget.get_value()))
+    switch_id = MockObj(config[CONF_ID], "->")
+    v = literal("v")
     async with LambdaContext([(cg.bool_, "v")]) as control:
-        with LvConditional(MockObj("v")) as cond:
+        with LvConditional(v) as cond:
             widget.add_state(LV_STATE.CHECKED)
             cond.else_()
             widget.clear_state(LV_STATE.CHECKED)
         lv.event_send(widget.obj, API_EVENT, cg.nullptr)
-        control.add(switch.publish_state(literal("v")))
+        control.add(switch_id.publish_state(v))
+    switch = cg.new_Pvariable(config[CONF_ID], await control.get_lambda())
+    await cg.register_component(switch, config)
+    await register_switch(switch, config)
+    async with LambdaContext(EVENT_ARG) as checked_ctx:
+        checked_ctx.add(switch.publish_state(widget.get_value()))
     async with LvContext() as ctx:
-        lv_add(switch.set_control_lambda(await control.get_lambda()))
         ctx.add(
             lvgl_static.add_event_cb(
                 widget.obj,
diff --git a/esphome/components/lvgl/switch/lvgl_switch.h b/esphome/components/lvgl/switch/lvgl_switch.h
index af839b8892..485459691c 100644
--- a/esphome/components/lvgl/switch/lvgl_switch.h
+++ b/esphome/components/lvgl/switch/lvgl_switch.h
@@ -10,26 +10,15 @@
 namespace esphome {
 namespace lvgl {
 
-class LVGLSwitch : public switch_::Switch {
+class LVGLSwitch : public switch_::Switch, public Component {
  public:
-  void set_control_lambda(std::function<void(bool)> state_lambda) {
-    this->state_lambda_ = std::move(state_lambda);
-    if (this->initial_state_.has_value()) {
-      this->state_lambda_(this->initial_state_.value());
-      this->initial_state_.reset();
-    }
-  }
+  LVGLSwitch(std::function<void(bool)> state_lambda) : state_lambda_(std::move(state_lambda)) {}
+
+  void setup() override { this->write_state(this->get_initial_state_with_restore_mode().value_or(false)); }
 
  protected:
-  void write_state(bool value) override {
-    if (this->state_lambda_ != nullptr) {
-      this->state_lambda_(value);
-    } else {
-      this->initial_state_ = value;
-    }
-  }
+  void write_state(bool value) override { this->state_lambda_(value); }
   std::function<void(bool)> state_lambda_{};
-  optional<bool> initial_state_{};
 };
 
 }  // namespace lvgl