From 15fe049a99143e55212532a88ee90f82b3a31f08 Mon Sep 17 00:00:00 2001
From: Joshua Spence <josh@joshuaspence.com>
Date: Mon, 10 Jan 2022 05:47:00 +1100
Subject: [PATCH] Add `restore_mode` to output switch (#3016)

Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
---
 esphome/components/output/switch/__init__.py  | 17 +++++++++-
 .../output/switch/output_switch.cpp           | 33 ++++++++++++++-----
 .../components/output/switch/output_switch.h  | 12 +++++++
 3 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/esphome/components/output/switch/__init__.py b/esphome/components/output/switch/__init__.py
index 11d073d28c..46135d117e 100644
--- a/esphome/components/output/switch/__init__.py
+++ b/esphome/components/output/switch/__init__.py
@@ -1,15 +1,28 @@
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.components import output, switch
-from esphome.const import CONF_ID, CONF_OUTPUT
+from esphome.const import CONF_ID, CONF_OUTPUT, CONF_RESTORE_MODE
 from .. import output_ns
 
 OutputSwitch = output_ns.class_("OutputSwitch", switch.Switch, cg.Component)
 
+OutputSwitchRestoreMode = output_ns.enum("OutputSwitchRestoreMode")
+RESTORE_MODES = {
+    "RESTORE_DEFAULT_OFF": OutputSwitchRestoreMode.OUTPUT_SWITCH_RESTORE_DEFAULT_OFF,
+    "RESTORE_DEFAULT_ON": OutputSwitchRestoreMode.OUTPUT_SWITCH_RESTORE_DEFAULT_ON,
+    "ALWAYS_OFF": OutputSwitchRestoreMode.OUTPUT_SWITCH_ALWAYS_OFF,
+    "ALWAYS_ON": OutputSwitchRestoreMode.OUTPUT_SWITCH_ALWAYS_ON,
+    "RESTORE_INVERTED_DEFAULT_OFF": OutputSwitchRestoreMode.OUTPUT_SWITCH_RESTORE_INVERTED_DEFAULT_OFF,
+    "RESTORE_INVERTED_DEFAULT_ON": OutputSwitchRestoreMode.OUTPUT_SWITCH_RESTORE_INVERTED_DEFAULT_ON,
+}
+
 CONFIG_SCHEMA = switch.SWITCH_SCHEMA.extend(
     {
         cv.GenerateID(): cv.declare_id(OutputSwitch),
         cv.Required(CONF_OUTPUT): cv.use_id(output.BinaryOutput),
+        cv.Optional(CONF_RESTORE_MODE, default="RESTORE_DEFAULT_OFF"): cv.enum(
+            RESTORE_MODES, upper=True, space="_"
+        ),
     }
 ).extend(cv.COMPONENT_SCHEMA)
 
@@ -21,3 +34,5 @@ async def to_code(config):
 
     output_ = await cg.get_variable(config[CONF_OUTPUT])
     cg.add(var.set_output(output_))
+
+    cg.add(var.set_restore_mode(config[CONF_RESTORE_MODE]))
diff --git a/esphome/components/output/switch/output_switch.cpp b/esphome/components/output/switch/output_switch.cpp
index 8db45f3a2b..3691896cbe 100644
--- a/esphome/components/output/switch/output_switch.cpp
+++ b/esphome/components/output/switch/output_switch.cpp
@@ -8,15 +8,32 @@ static const char *const TAG = "output.switch";
 
 void OutputSwitch::dump_config() { LOG_SWITCH("", "Output Switch", this); }
 void OutputSwitch::setup() {
-  auto restored = this->get_initial_state();
-  if (!restored.has_value())
-    return;
-
-  if (*restored) {
-    this->turn_on();
-  } else {
-    this->turn_off();
+  bool initial_state = false;
+  switch (this->restore_mode_) {
+    case OUTPUT_SWITCH_RESTORE_DEFAULT_OFF:
+      initial_state = this->get_initial_state().value_or(false);
+      break;
+    case OUTPUT_SWITCH_RESTORE_DEFAULT_ON:
+      initial_state = this->get_initial_state().value_or(true);
+      break;
+    case OUTPUT_SWITCH_RESTORE_INVERTED_DEFAULT_OFF:
+      initial_state = !this->get_initial_state().value_or(true);
+      break;
+    case OUTPUT_SWITCH_RESTORE_INVERTED_DEFAULT_ON:
+      initial_state = !this->get_initial_state().value_or(false);
+      break;
+    case OUTPUT_SWITCH_ALWAYS_OFF:
+      initial_state = false;
+      break;
+    case OUTPUT_SWITCH_ALWAYS_ON:
+      initial_state = true;
+      break;
   }
+
+  if (initial_state)
+    this->turn_on();
+  else
+    this->turn_off();
 }
 void OutputSwitch::write_state(bool state) {
   if (state) {
diff --git a/esphome/components/output/switch/output_switch.h b/esphome/components/output/switch/output_switch.h
index a184a342fe..fc2c110662 100644
--- a/esphome/components/output/switch/output_switch.h
+++ b/esphome/components/output/switch/output_switch.h
@@ -7,10 +7,21 @@
 namespace esphome {
 namespace output {
 
+enum OutputSwitchRestoreMode {
+  OUTPUT_SWITCH_RESTORE_DEFAULT_OFF,
+  OUTPUT_SWITCH_RESTORE_DEFAULT_ON,
+  OUTPUT_SWITCH_ALWAYS_OFF,
+  OUTPUT_SWITCH_ALWAYS_ON,
+  OUTPUT_SWITCH_RESTORE_INVERTED_DEFAULT_OFF,
+  OUTPUT_SWITCH_RESTORE_INVERTED_DEFAULT_ON,
+};
+
 class OutputSwitch : public switch_::Switch, public Component {
  public:
   void set_output(BinaryOutput *output) { output_ = output; }
 
+  void set_restore_mode(OutputSwitchRestoreMode restore_mode) { restore_mode_ = restore_mode; }
+
   void setup() override;
   float get_setup_priority() const override { return setup_priority::HARDWARE - 1.0f; }
   void dump_config() override;
@@ -19,6 +30,7 @@ class OutputSwitch : public switch_::Switch, public Component {
   void write_state(bool state) override;
 
   output::BinaryOutput *output_;
+  OutputSwitchRestoreMode restore_mode_;
 };
 
 }  // namespace output