From 8e29437900927607b607d9ce235444e8cdcd3f84 Mon Sep 17 00:00:00 2001
From: Samuel Sieb <samuel-github@sieb.net>
Date: Thu, 8 May 2025 01:26:10 -0700
Subject: [PATCH] [key_collector] enable/disable (#8718)

Co-authored-by: Samuel Sieb <samuel@sieb.net>
---
 esphome/components/key_collector/__init__.py  | 35 +++++++++++++++++++
 .../key_collector/key_collector.cpp           |  8 +++++
 .../components/key_collector/key_collector.h  | 10 ++++++
 tests/components/key_collector/common.yaml    |  8 +++++
 4 files changed, 61 insertions(+)

diff --git a/esphome/components/key_collector/__init__.py b/esphome/components/key_collector/__init__.py
index 5750812f5c..17af40da1a 100644
--- a/esphome/components/key_collector/__init__.py
+++ b/esphome/components/key_collector/__init__.py
@@ -3,6 +3,7 @@ import esphome.codegen as cg
 from esphome.components import key_provider
 import esphome.config_validation as cv
 from esphome.const import (
+    CONF_ENABLE_ON_BOOT,
     CONF_ID,
     CONF_MAX_LENGTH,
     CONF_MIN_LENGTH,
@@ -28,6 +29,8 @@ CONF_ON_RESULT = "on_result"
 
 key_collector_ns = cg.esphome_ns.namespace("key_collector")
 KeyCollector = key_collector_ns.class_("KeyCollector", cg.Component)
+EnableAction = key_collector_ns.class_("EnableAction", automation.Action)
+DisableAction = key_collector_ns.class_("DisableAction", automation.Action)
 
 CONFIG_SCHEMA = cv.All(
     cv.COMPONENT_SCHEMA.extend(
@@ -46,6 +49,7 @@ CONFIG_SCHEMA = cv.All(
             cv.Optional(CONF_ON_RESULT): automation.validate_automation(single=True),
             cv.Optional(CONF_ON_TIMEOUT): automation.validate_automation(single=True),
             cv.Optional(CONF_TIMEOUT): cv.positive_time_period_milliseconds,
+            cv.Optional(CONF_ENABLE_ON_BOOT, default=True): cv.boolean,
         }
     ),
     cv.has_at_least_one_key(CONF_END_KEYS, CONF_MAX_LENGTH),
@@ -94,3 +98,34 @@ async def to_code(config):
         )
     if CONF_TIMEOUT in config:
         cg.add(var.set_timeout(config[CONF_TIMEOUT]))
+    cg.add(var.set_enabled(config[CONF_ENABLE_ON_BOOT]))
+
+
+@automation.register_action(
+    "key_collector.enable",
+    EnableAction,
+    automation.maybe_simple_id(
+        {
+            cv.GenerateID(): cv.use_id(KeyCollector),
+        }
+    ),
+)
+async def enable_to_code(config, action_id, template_arg, args):
+    var = cg.new_Pvariable(action_id, template_arg)
+    await cg.register_parented(var, config[CONF_ID])
+    return var
+
+
+@automation.register_action(
+    "key_collector.disable",
+    DisableAction,
+    automation.maybe_simple_id(
+        {
+            cv.GenerateID(): cv.use_id(KeyCollector),
+        }
+    ),
+)
+async def disable_to_code(config, action_id, template_arg, args):
+    var = cg.new_Pvariable(action_id, template_arg)
+    await cg.register_parented(var, config[CONF_ID])
+    return var
diff --git a/esphome/components/key_collector/key_collector.cpp b/esphome/components/key_collector/key_collector.cpp
index bf2333d97d..ffb4b47fa2 100644
--- a/esphome/components/key_collector/key_collector.cpp
+++ b/esphome/components/key_collector/key_collector.cpp
@@ -45,6 +45,12 @@ void KeyCollector::set_provider(key_provider::KeyProvider *provider) {
   provider->add_on_key_callback([this](uint8_t key) { this->key_pressed_(key); });
 }
 
+void KeyCollector::set_enabled(bool enabled) {
+  this->enabled_ = enabled;
+  if (!enabled)
+    this->clear(false);
+}
+
 void KeyCollector::clear(bool progress_update) {
   this->result_.clear();
   this->start_key_ = 0;
@@ -55,6 +61,8 @@ void KeyCollector::clear(bool progress_update) {
 void KeyCollector::send_key(uint8_t key) { this->key_pressed_(key); }
 
 void KeyCollector::key_pressed_(uint8_t key) {
+  if (!this->enabled_)
+    return;
   this->last_key_time_ = millis();
   if (!this->start_keys_.empty() && !this->start_key_) {
     if (this->start_keys_.find(key) != std::string::npos) {
diff --git a/esphome/components/key_collector/key_collector.h b/esphome/components/key_collector/key_collector.h
index 7ef53929ef..6e585ddd8e 100644
--- a/esphome/components/key_collector/key_collector.h
+++ b/esphome/components/key_collector/key_collector.h
@@ -25,6 +25,7 @@ class KeyCollector : public Component {
   Trigger<std::string, uint8_t, uint8_t> *get_result_trigger() const { return this->result_trigger_; };
   Trigger<std::string, uint8_t> *get_timeout_trigger() const { return this->timeout_trigger_; };
   void set_timeout(int timeout) { this->timeout_ = timeout; };
+  void set_enabled(bool enabled);
 
   void clear(bool progress_update = true);
   void send_key(uint8_t key);
@@ -47,6 +48,15 @@ class KeyCollector : public Component {
   Trigger<std::string, uint8_t> *timeout_trigger_;
   uint32_t last_key_time_;
   uint32_t timeout_{0};
+  bool enabled_;
+};
+
+template<typename... Ts> class EnableAction : public Action<Ts...>, public Parented<KeyCollector> {
+  void play(Ts... x) override { this->parent_->set_enabled(true); }
+};
+
+template<typename... Ts> class DisableAction : public Action<Ts...>, public Parented<KeyCollector> {
+  void play(Ts... x) override { this->parent_->set_enabled(false); }
 };
 
 }  // namespace key_collector
diff --git a/tests/components/key_collector/common.yaml b/tests/components/key_collector/common.yaml
index d58922ca91..12e541c865 100644
--- a/tests/components/key_collector/common.yaml
+++ b/tests/components/key_collector/common.yaml
@@ -26,3 +26,11 @@ key_collector:
       - logger.log:
           format: "input timeout: '%s', started by '%c'"
           args: ['x.c_str()', "(start == 0 ? '~' : start)"]
+    enable_on_boot: false
+
+button:
+  - platform: template
+    id: button0
+    on_press:
+      - key_collector.enable:
+      - key_collector.disable: