refactor restore_state_ to eliminate awkward setup delegation

esphome/components/template/water_heater/__init__.py

@@ -20,7 +20,7 @@ from .. import template_ns
 CONF_CURRENT_TEMPERATURE = "current_temperature"
 
 TemplateWaterHeater = template_ns.class_(
-    "TemplateWaterHeater", water_heater.WaterHeater
+    "TemplateWaterHeater", water_heater.WaterHeater, cg.Component
 )
 
 TemplateWaterHeaterPublishAction = template_ns.class_(
@@ -36,24 +36,29 @@ RESTORE_MODES = {
     "RESTORE_AND_CALL": TemplateWaterHeaterRestoreMode.WATER_HEATER_RESTORE_AND_CALL,
 }
 
-CONFIG_SCHEMA = water_heater.water_heater_schema(TemplateWaterHeater).extend(
-    {
-        cv.Optional(CONF_OPTIMISTIC, default=True): cv.boolean,
-        cv.Optional(CONF_SET_ACTION): automation.validate_automation(single=True),
-        cv.Optional(CONF_RESTORE_MODE, default="NO_RESTORE"): cv.enum(
-            RESTORE_MODES, upper=True
-        ),
-        cv.Optional(CONF_CURRENT_TEMPERATURE): cv.returning_lambda,
-        cv.Optional(CONF_MODE): cv.returning_lambda,
-        cv.Optional(CONF_SUPPORTED_MODES): cv.ensure_list(
-            water_heater.validate_water_heater_mode
-        ),
-    }
+CONFIG_SCHEMA = (
+    water_heater.water_heater_schema(TemplateWaterHeater)
+    .extend(
+        {
+            cv.Optional(CONF_OPTIMISTIC, default=True): cv.boolean,
+            cv.Optional(CONF_SET_ACTION): automation.validate_automation(single=True),
+            cv.Optional(CONF_RESTORE_MODE, default="NO_RESTORE"): cv.enum(
+                RESTORE_MODES, upper=True
+            ),
+            cv.Optional(CONF_CURRENT_TEMPERATURE): cv.returning_lambda,
+            cv.Optional(CONF_MODE): cv.returning_lambda,
+            cv.Optional(CONF_SUPPORTED_MODES): cv.ensure_list(
+                water_heater.validate_water_heater_mode
+            ),
+        }
+    )
+    .extend(cv.COMPONENT_SCHEMA)
 )
 
 
 async def to_code(config: ConfigType) -> None:
     var = cg.new_Pvariable(config[CONF_ID])
+    await cg.register_component(var, config)
     await water_heater.register_water_heater(var, config)
 
     cg.add(var.set_optimistic(config[CONF_OPTIMISTIC]))

esphome/components/template/water_heater/template_water_heater.cpp

@@ -10,7 +10,7 @@ TemplateWaterHeater::TemplateWaterHeater() : set_trigger_(new Trigger<>()) {}
 void TemplateWaterHeater::setup() {
   if (this->restore_mode_ == TemplateWaterHeaterRestoreMode::WATER_HEATER_RESTORE ||
       this->restore_mode_ == TemplateWaterHeaterRestoreMode::WATER_HEATER_RESTORE_AND_CALL) {
-    auto restore = this->restore_state();
+    auto restore = this->restore_state_();
 
     if (restore.has_value()) {
       restore->perform();

esphome/components/template/water_heater/template_water_heater.h

@@ -13,7 +13,7 @@ enum TemplateWaterHeaterRestoreMode {
   WATER_HEATER_RESTORE_AND_CALL,
 };
 
-class TemplateWaterHeater : public water_heater::WaterHeater {
+class TemplateWaterHeater : public Component, public water_heater::WaterHeater {
  public:
   TemplateWaterHeater();
 

esphome/components/water_heater/__init__.py

@@ -18,7 +18,7 @@ CODEOWNERS = ["@dhoeben"]
 IS_PLATFORM_COMPONENT = True
 
 water_heater_ns = cg.esphome_ns.namespace("water_heater")
-WaterHeater = water_heater_ns.class_("WaterHeater", cg.EntityBase, cg.Component)
+WaterHeater = water_heater_ns.class_("WaterHeater", cg.EntityBase)
 WaterHeaterCall = water_heater_ns.class_("WaterHeaterCall")
 WaterHeaterTraits = water_heater_ns.class_("WaterHeaterTraits")
 
@@ -46,7 +46,7 @@ _WATER_HEATER_SCHEMA = cv.ENTITY_BASE_SCHEMA.extend(
             }
         ),
     }
-).extend(cv.COMPONENT_SCHEMA)
+)
 
 _WATER_HEATER_SCHEMA.add_extra(entity_duplicate_validator("water_heater"))
 
@@ -91,8 +91,6 @@ async def register_water_heater(var: cg.Pvariable, config: ConfigType) -> cg.Pva
 
     cg.add_define("USE_WATER_HEATER")
 
-    await cg.register_component(var, config)
-
     cg.add(cg.App.register_water_heater(var))
 
     CORE.register_platform_component("water_heater", var)

esphome/components/water_heater/water_heater.cpp

@@ -146,10 +146,6 @@ void WaterHeaterCall::validate_() {
   }
 }
 
-void WaterHeater::setup() {
-  this->pref_ = global_preferences->make_preference<SavedWaterHeaterState>(this->get_preference_hash());
-}
-
 void WaterHeater::publish_state() {
   auto traits = this->get_traits();
   ESP_LOGD(TAG,
@@ -188,7 +184,8 @@ void WaterHeater::publish_state() {
   this->pref_.save(&saved);
 }
 
-optional<WaterHeaterCall> WaterHeater::restore_state() {
+optional<WaterHeaterCall> WaterHeater::restore_state_() {
+  this->pref_ = global_preferences->make_preference<SavedWaterHeaterState>(this->get_preference_hash());
   SavedWaterHeaterState recovered{};
   if (!this->pref_.load(&recovered))
     return {};

esphome/components/water_heater/water_heater.h

@@ -177,7 +177,7 @@ class WaterHeaterTraits {
   WaterHeaterModeMask supported_modes_;
 };
 
-class WaterHeater : public EntityBase, public Component {
+class WaterHeater : public EntityBase {
  public:
   WaterHeaterMode get_mode() const { return this->mode_; }
   float get_current_temperature() const { return this->current_temperature_; }
@@ -204,16 +204,15 @@ class WaterHeater : public EntityBase, public Component {
 #endif
   virtual void control(const WaterHeaterCall &call) = 0;
 
-  void setup() override;
-
-  optional<WaterHeaterCall> restore_state();
-
  protected:
   virtual WaterHeaterTraits traits() = 0;
 
   /// Log the traits of this water heater for dump_config().
   void dump_traits_(const char *tag);
 
+  /// Restore the state of the water heater, call this from your setup() method.
+  optional<WaterHeaterCall> restore_state_();
+
   /// Set the mode of the water heater. Should only be called from control().
   void set_mode_(WaterHeaterMode mode) { this->mode_ = mode; }
   /// Set the target temperature of the water heater. Should only be called from control().