From 268387f829d0965113bc286594e335e0337cd7be Mon Sep 17 00:00:00 2001
From: Otto Winter <otto@otto-winter.com>
Date: Sun, 20 Oct 2019 16:15:30 +0200
Subject: [PATCH] Add script.wait action (#778)

Fixes https://github.com/esphome/feature-requests/issues/416, fixes https://github.com/esphome/issues/issues/572
---
 esphome/components/script/__init__.py |  9 ++++++
 esphome/components/script/script.h    | 42 +++++++++++++++++++++++++++
 tests/test3.yaml                      |  2 ++
 3 files changed, 53 insertions(+)

diff --git a/esphome/components/script/__init__.py b/esphome/components/script/__init__.py
index e1983689a6..9590679f83 100644
--- a/esphome/components/script/__init__.py
+++ b/esphome/components/script/__init__.py
@@ -8,6 +8,7 @@ script_ns = cg.esphome_ns.namespace('script')
 Script = script_ns.class_('Script', automation.Trigger.template())
 ScriptExecuteAction = script_ns.class_('ScriptExecuteAction', automation.Action)
 ScriptStopAction = script_ns.class_('ScriptStopAction', automation.Action)
+ScriptWaitAction = script_ns.class_('ScriptWaitAction', automation.Action)
 IsRunningCondition = script_ns.class_('IsRunningCondition', automation.Condition)
 
 CONFIG_SCHEMA = automation.validate_automation({
@@ -42,6 +43,14 @@ def script_stop_action_to_code(config, action_id, template_arg, args):
     yield cg.new_Pvariable(action_id, template_arg, paren)
 
 
+@automation.register_action('script.wait', ScriptWaitAction, maybe_simple_id({
+    cv.Required(CONF_ID): cv.use_id(Script)
+}))
+def script_wait_action_to_code(config, action_id, template_arg, args):
+    paren = yield cg.get_variable(config[CONF_ID])
+    yield cg.new_Pvariable(action_id, template_arg, paren)
+
+
 @automation.register_condition('script.is_running', IsRunningCondition, automation.maybe_simple_id({
     cv.Required(CONF_ID): cv.use_id(Script)
 }))
diff --git a/esphome/components/script/script.h b/esphome/components/script/script.h
index f937b9d637..3b97327da8 100644
--- a/esphome/components/script/script.h
+++ b/esphome/components/script/script.h
@@ -49,5 +49,47 @@ template<typename... Ts> class IsRunningCondition : public Condition<Ts...> {
   Script *parent_;
 };
 
+template<typename... Ts> class ScriptWaitAction : public Action<Ts...>, public Component {
+ public:
+  ScriptWaitAction(Script *script) : script_(script) {}
+
+  void play(Ts... x) { /* ignore - see play_complex */
+  }
+
+  void play_complex(Ts... x) override {
+    // Check if we can continue immediately.
+    if (!this->script_->is_running()) {
+      this->triggered_ = false;
+      this->play_next(x...);
+      return;
+    }
+    this->var_ = std::make_tuple(x...);
+    this->triggered_ = true;
+    this->loop();
+  }
+
+  void stop() override { this->triggered_ = false; }
+
+  void loop() override {
+    if (!this->triggered_)
+      return;
+
+    if (this->script_->is_running())
+      return;
+
+    this->triggered_ = false;
+    this->play_next_tuple(this->var_);
+  }
+
+  float get_setup_priority() const override { return setup_priority::DATA; }
+
+  bool is_running() override { return this->triggered_ || this->is_running_next(); }
+
+ protected:
+  Script *script_;
+  bool triggered_{false};
+  std::tuple<Ts...> var_{};
+};
+
 }  // namespace script
 }  // namespace esphome
diff --git a/tests/test3.yaml b/tests/test3.yaml
index df4afe5ad9..f3e530f07a 100644
--- a/tests/test3.yaml
+++ b/tests/test3.yaml
@@ -383,6 +383,8 @@ text_sensor:
       - lambda: !lambda |-
           ESP_LOGD("main", "The state is %s=%s", x.c_str(), id(version_sensor).state.c_str());
       - script.execute: my_script
+      - script.wait: my_script
+      - script.stop: my_script
       - homeassistant.service:
           service: notify.html5
           data: