[water_heater] (2/4) Implement template for new water_heater component (#12516)

Co-authored-by: J. Nick Koston <nick@koston.org>
Co-authored-by: J. Nick Koston <nick@home-assistant.io>

esphome/components/template/water_heater/__init__.py

@@ -0,0 +1,123 @@
+from esphome import automation
+import esphome.codegen as cg
+from esphome.components import water_heater
+import esphome.config_validation as cv
+from esphome.const import (
+    CONF_ID,
+    CONF_MODE,
+    CONF_OPTIMISTIC,
+    CONF_RESTORE_MODE,
+    CONF_SET_ACTION,
+    CONF_SUPPORTED_MODES,
+    CONF_TARGET_TEMPERATURE,
+)
+from esphome.core import ID
+from esphome.cpp_generator import MockObj, TemplateArgsType
+from esphome.types import ConfigType
+
+from .. import template_ns
+
+CONF_CURRENT_TEMPERATURE = "current_temperature"
+
+TemplateWaterHeater = template_ns.class_(
+    "TemplateWaterHeater", water_heater.WaterHeater
+)
+
+TemplateWaterHeaterPublishAction = template_ns.class_(
+    "TemplateWaterHeaterPublishAction",
+    automation.Action,
+    cg.Parented.template(TemplateWaterHeater),
+)
+
+TemplateWaterHeaterRestoreMode = template_ns.enum("TemplateWaterHeaterRestoreMode")
+RESTORE_MODES = {
+    "NO_RESTORE": TemplateWaterHeaterRestoreMode.WATER_HEATER_NO_RESTORE,
+    "RESTORE": TemplateWaterHeaterRestoreMode.WATER_HEATER_RESTORE,
+    "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
+        ),
+    }
+)
+
+
+async def to_code(config: ConfigType) -> None:
+    var = cg.new_Pvariable(config[CONF_ID])
+    await water_heater.register_water_heater(var, config)
+
+    cg.add(var.set_optimistic(config[CONF_OPTIMISTIC]))
+
+    if CONF_SET_ACTION in config:
+        await automation.build_automation(
+            var.get_set_trigger(), [], config[CONF_SET_ACTION]
+        )
+
+    cg.add(var.set_restore_mode(config[CONF_RESTORE_MODE]))
+
+    if CONF_CURRENT_TEMPERATURE in config:
+        template_ = await cg.process_lambda(
+            config[CONF_CURRENT_TEMPERATURE],
+            [],
+            return_type=cg.optional.template(cg.float_),
+        )
+        cg.add(var.set_current_temperature_lambda(template_))
+
+    if CONF_MODE in config:
+        template_ = await cg.process_lambda(
+            config[CONF_MODE],
+            [],
+            return_type=cg.optional.template(water_heater.WaterHeaterMode),
+        )
+        cg.add(var.set_mode_lambda(template_))
+
+    if CONF_SUPPORTED_MODES in config:
+        cg.add(var.set_supported_modes(config[CONF_SUPPORTED_MODES]))
+
+
+@automation.register_action(
+    "water_heater.template.publish",
+    TemplateWaterHeaterPublishAction,
+    cv.Schema(
+        {
+            cv.GenerateID(): cv.use_id(TemplateWaterHeater),
+            cv.Optional(CONF_CURRENT_TEMPERATURE): cv.templatable(cv.temperature),
+            cv.Optional(CONF_TARGET_TEMPERATURE): cv.templatable(cv.temperature),
+            cv.Optional(CONF_MODE): cv.templatable(
+                water_heater.validate_water_heater_mode
+            ),
+        }
+    ),
+)
+async def water_heater_template_publish_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
+    var = cg.new_Pvariable(action_id, template_arg)
+    await cg.register_parented(var, config[CONF_ID])
+
+    if current_temp := config.get(CONF_CURRENT_TEMPERATURE):
+        template_ = await cg.templatable(current_temp, args, float)
+        cg.add(var.set_current_temperature(template_))
+
+    if target_temp := config.get(CONF_TARGET_TEMPERATURE):
+        template_ = await cg.templatable(target_temp, args, float)
+        cg.add(var.set_target_temperature(template_))
+
+    if mode := config.get(CONF_MODE):
+        template_ = await cg.templatable(mode, args, water_heater.WaterHeaterMode)
+        cg.add(var.set_mode(template_))
+
+    return var

esphome/components/template/water_heater/automation.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "template_water_heater.h"
+#include "esphome/core/automation.h"
+
+namespace esphome::template_ {
+
+template<typename... Ts>
+class TemplateWaterHeaterPublishAction : public Action<Ts...>, public Parented<TemplateWaterHeater> {
+ public:
+  TEMPLATABLE_VALUE(float, current_temperature)
+  TEMPLATABLE_VALUE(float, target_temperature)
+  TEMPLATABLE_VALUE(water_heater::WaterHeaterMode, mode)
+
+  void play(const Ts &...x) override {
+    if (this->current_temperature_.has_value()) {
+      this->parent_->set_current_temperature(this->current_temperature_.value(x...));
+    }
+    bool needs_call = this->target_temperature_.has_value() || this->mode_.has_value();
+    if (needs_call) {
+      auto call = this->parent_->make_call();
+      if (this->target_temperature_.has_value()) {
+        call.set_target_temperature(this->target_temperature_.value(x...));
+      }
+      if (this->mode_.has_value()) {
+        call.set_mode(this->mode_.value(x...));
+      }
+      call.perform();
+    } else {
+      this->parent_->publish_state();
+    }
+  }
+};
+
+}  // namespace esphome::template_

esphome/components/template/water_heater/template_water_heater.cpp

@@ -0,0 +1,88 @@
+#include "template_water_heater.h"
+#include "esphome/core/log.h"
+
+namespace esphome::template_ {
+
+static const char *const TAG = "template.water_heater";
+
+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();
+
+    if (restore.has_value()) {
+      restore->perform();
+    }
+  }
+  if (!this->current_temperature_f_.has_value() && !this->mode_f_.has_value())
+    this->disable_loop();
+}
+
+water_heater::WaterHeaterTraits TemplateWaterHeater::traits() {
+  auto traits = water_heater::WaterHeater::get_traits();
+
+  if (!this->supported_modes_.empty()) {
+    traits.set_supported_modes(this->supported_modes_);
+  }
+
+  traits.set_supports_current_temperature(true);
+  return traits;
+}
+
+void TemplateWaterHeater::loop() {
+  bool changed = false;
+
+  auto curr_temp = this->current_temperature_f_.call();
+  if (curr_temp.has_value()) {
+    if (*curr_temp != this->current_temperature_) {
+      this->current_temperature_ = *curr_temp;
+      changed = true;
+    }
+  }
+
+  auto new_mode = this->mode_f_.call();
+  if (new_mode.has_value()) {
+    if (*new_mode != this->mode_) {
+      this->mode_ = *new_mode;
+      changed = true;
+    }
+  }
+
+  if (changed) {
+    this->publish_state();
+  }
+}
+
+void TemplateWaterHeater::dump_config() {
+  LOG_WATER_HEATER("", "Template Water Heater", this);
+  ESP_LOGCONFIG(TAG, "  Optimistic: %s", YESNO(this->optimistic_));
+}
+
+float TemplateWaterHeater::get_setup_priority() const { return setup_priority::HARDWARE; }
+
+water_heater::WaterHeaterCallInternal TemplateWaterHeater::make_call() {
+  return water_heater::WaterHeaterCallInternal(this);
+}
+
+void TemplateWaterHeater::control(const water_heater::WaterHeaterCall &call) {
+  if (call.get_mode().has_value()) {
+    if (this->optimistic_) {
+      this->mode_ = *call.get_mode();
+    }
+  }
+  if (!std::isnan(call.get_target_temperature())) {
+    if (this->optimistic_) {
+      this->target_temperature_ = call.get_target_temperature();
+    }
+  }
+
+  this->set_trigger_->trigger();
+
+  if (this->optimistic_) {
+    this->publish_state();
+  }
+}
+
+}  // namespace esphome::template_

esphome/components/template/water_heater/template_water_heater.h

@@ -0,0 +1,53 @@
+#pragma once
+
+#include "esphome/core/component.h"
+#include "esphome/core/automation.h"
+#include "esphome/core/template_lambda.h"
+#include "esphome/components/water_heater/water_heater.h"
+
+namespace esphome::template_ {
+
+enum TemplateWaterHeaterRestoreMode {
+  WATER_HEATER_NO_RESTORE,
+  WATER_HEATER_RESTORE,
+  WATER_HEATER_RESTORE_AND_CALL,
+};
+
+class TemplateWaterHeater : public water_heater::WaterHeater {
+ public:
+  TemplateWaterHeater();
+
+  template<typename F> void set_current_temperature_lambda(F &&f) {
+    this->current_temperature_f_.set(std::forward<F>(f));
+  }
+  template<typename F> void set_mode_lambda(F &&f) { this->mode_f_.set(std::forward<F>(f)); }
+
+  void set_optimistic(bool optimistic) { this->optimistic_ = optimistic; }
+  void set_restore_mode(TemplateWaterHeaterRestoreMode restore_mode) { this->restore_mode_ = restore_mode; }
+  void set_supported_modes(const std::initializer_list<water_heater::WaterHeaterMode> &modes) {
+    this->supported_modes_ = modes;
+  }
+
+  Trigger<> *get_set_trigger() const { return this->set_trigger_; }
+
+  void setup() override;
+  void loop() override;
+  void dump_config() override;
+  float get_setup_priority() const override;
+
+  water_heater::WaterHeaterCallInternal make_call() override;
+
+ protected:
+  void control(const water_heater::WaterHeaterCall &call) override;
+  water_heater::WaterHeaterTraits traits() override;
+
+  // Ordered to minimize padding on 32-bit: 4-byte members first, then smaller
+  Trigger<> *set_trigger_;
+  TemplateLambda<float> current_temperature_f_;
+  TemplateLambda<water_heater::WaterHeaterMode> mode_f_;
+  TemplateWaterHeaterRestoreMode restore_mode_{WATER_HEATER_NO_RESTORE};
+  water_heater::WaterHeaterModeMask supported_modes_;
+  bool optimistic_{true};
+};
+
+}  // namespace esphome::template_

tests/components/template/common-base.yaml

@@ -9,6 +9,18 @@ esphome:
         id: template_sens
         state: !lambda "return 42.0;"
 
+    - water_heater.template.publish:
+        id: template_water_heater
+        target_temperature: 50.0
+        mode: ECO
+
+    # Templated
+    - water_heater.template.publish:
+        id: template_water_heater
+        current_temperature: !lambda "return 45.0;"
+        target_temperature: !lambda "return 55.0;"
+        mode: !lambda "return water_heater::WATER_HEATER_MODE_GAS;"
+
     # Test C++ API: set_template() with stateless lambda (no captures)
     # NOTE: set_template() is not intended to be a public API, but we test it to ensure it doesn't break.
     - lambda: |-
@@ -299,6 +311,24 @@ alarm_control_panel:
     codes:
       - "1234"
 
+water_heater:
+  - platform: template
+    id: template_water_heater
+    name: "Template Water Heater"
+    optimistic: true
+    current_temperature: !lambda "return 42.0f;"
+    mode: !lambda "return water_heater::WATER_HEATER_MODE_ECO;"
+    supported_modes:
+      - "OFF"
+      - ECO
+      - GAS
+      - ELECTRIC
+      - HEAT_PUMP
+      - HIGH_DEMAND
+      - PERFORMANCE
+    set_action:
+      - logger.log: "set_action"
+
 datetime:
   - platform: template
     name: Date