From d8e33c5a69aa0053ae01352341e57265df7df802 Mon Sep 17 00:00:00 2001
From: Oxan van Leeuwen <oxan@oxanvanleeuwen.nl>
Date: Wed, 10 Nov 2021 19:30:07 +0100
Subject: [PATCH] Add repeat action for automations (#2538)

Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
---
 esphome/automation.py          | 21 +++++++++++++++++++++
 esphome/core/base_automation.h | 33 +++++++++++++++++++++++++++++++++
 tests/test5.yaml               |  8 ++++++++
 3 files changed, 62 insertions(+)

diff --git a/esphome/automation.py b/esphome/automation.py
index 0768bf8869..fab998527f 100644
--- a/esphome/automation.py
+++ b/esphome/automation.py
@@ -3,6 +3,7 @@ import esphome.config_validation as cv
 from esphome.const import (
     CONF_AUTOMATION_ID,
     CONF_CONDITION,
+    CONF_COUNT,
     CONF_ELSE,
     CONF_ID,
     CONF_THEN,
@@ -66,6 +67,7 @@ DelayAction = cg.esphome_ns.class_("DelayAction", Action, cg.Component)
 LambdaAction = cg.esphome_ns.class_("LambdaAction", Action)
 IfAction = cg.esphome_ns.class_("IfAction", Action)
 WhileAction = cg.esphome_ns.class_("WhileAction", Action)
+RepeatAction = cg.esphome_ns.class_("RepeatAction", Action)
 WaitUntilAction = cg.esphome_ns.class_("WaitUntilAction", Action, cg.Component)
 UpdateComponentAction = cg.esphome_ns.class_("UpdateComponentAction", Action)
 Automation = cg.esphome_ns.class_("Automation")
@@ -241,6 +243,25 @@ async def while_action_to_code(config, action_id, template_arg, args):
     return var
 
 
+@register_action(
+    "repeat",
+    RepeatAction,
+    cv.Schema(
+        {
+            cv.Required(CONF_COUNT): cv.templatable(cv.positive_not_null_int),
+            cv.Required(CONF_THEN): validate_action_list,
+        }
+    ),
+)
+async def repeat_action_to_code(config, action_id, template_arg, args):
+    var = cg.new_Pvariable(action_id, template_arg)
+    count_template = await cg.templatable(config[CONF_COUNT], args, cg.uint32)
+    cg.add(var.set_count(count_template))
+    actions = await build_action_list(config[CONF_THEN], template_arg, args)
+    cg.add(var.add_then(actions))
+    return var
+
+
 def validate_wait_until(value):
     schema = cv.Schema(
         {
diff --git a/esphome/core/base_automation.h b/esphome/core/base_automation.h
index d97d369d33..e87a4a2765 100644
--- a/esphome/core/base_automation.h
+++ b/esphome/core/base_automation.h
@@ -224,6 +224,39 @@ template<typename... Ts> class WhileAction : public Action<Ts...> {
   std::tuple<Ts...> var_{};
 };
 
+template<typename... Ts> class RepeatAction : public Action<Ts...> {
+ public:
+  TEMPLATABLE_VALUE(uint32_t, count)
+
+  void add_then(const std::vector<Action<Ts...> *> &actions) {
+    this->then_.add_actions(actions);
+    this->then_.add_action(new LambdaAction<Ts...>([this](Ts... x) {
+      this->iteration_++;
+      if (this->iteration_ == this->count_.value(x...))
+        this->play_next_tuple_(this->var_);
+      else
+        this->then_.play_tuple(this->var_);
+    }));
+  }
+
+  void play_complex(Ts... x) override {
+    this->num_running_++;
+    this->var_ = std::make_tuple(x...);
+    this->iteration_ = 0;
+    this->then_.play_tuple(this->var_);
+  }
+
+  void play(Ts... x) override { /* ignore - see play_complex */
+  }
+
+  void stop() override { this->then_.stop(); }
+
+ protected:
+  uint32_t iteration_;
+  ActionList<Ts...> then_;
+  std::tuple<Ts...> var_;
+};
+
 template<typename... Ts> class WaitUntilAction : public Action<Ts...>, public Component {
  public:
   WaitUntilAction(Condition<Ts...> *condition) : condition_(condition) {}
diff --git a/tests/test5.yaml b/tests/test5.yaml
index 72df3ed212..f1fb786fe5 100644
--- a/tests/test5.yaml
+++ b/tests/test5.yaml
@@ -173,3 +173,11 @@ sensor:
     uart_id: uart2
     co2:
       name: CO2 Sensor
+
+script:
+  - id: automation_test
+    then:
+      - repeat:
+          count: 5
+          then:
+            - logger.log: "looping!"