From 04c12823b5187f40ed99812f715a2e853b1a54ca Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Mon, 20 Feb 2023 10:13:37 +1300
Subject: [PATCH] Simplify button_schema function (#4468)

---
 esphome/components/button/__init__.py         | 35 +++++++------------
 esphome/components/output/button/__init__.py  | 17 +++++----
 .../components/template/button/__init__.py    |  7 +---
 esphome/components/wake_on_lan/button.py      |  5 +--
 4 files changed, 26 insertions(+), 38 deletions(-)

diff --git a/esphome/components/button/__init__.py b/esphome/components/button/__init__.py
index 9d250bbc8b..55f2fe794a 100644
--- a/esphome/components/button/__init__.py
+++ b/esphome/components/button/__init__.py
@@ -57,34 +57,23 @@ _UNDEF = object()
 
 
 def button_schema(
-    class_: MockObjClass = _UNDEF,
+    class_: MockObjClass,
     *,
     icon: str = _UNDEF,
     entity_category: str = _UNDEF,
     device_class: str = _UNDEF,
 ) -> cv.Schema:
-    schema = BUTTON_SCHEMA
-    if class_ is not _UNDEF:
-        schema = schema.extend({cv.GenerateID(): cv.declare_id(class_)})
-    if icon is not _UNDEF:
-        schema = schema.extend({cv.Optional(CONF_ICON, default=icon): cv.icon})
-    if entity_category is not _UNDEF:
-        schema = schema.extend(
-            {
-                cv.Optional(
-                    CONF_ENTITY_CATEGORY, default=entity_category
-                ): cv.entity_category
-            }
-        )
-    if device_class is not _UNDEF:
-        schema = schema.extend(
-            {
-                cv.Optional(
-                    CONF_DEVICE_CLASS, default=device_class
-                ): validate_device_class
-            }
-        )
-    return schema
+    schema = {cv.GenerateID(): cv.declare_id(class_)}
+
+    for key, default, validator in [
+        (CONF_ICON, icon, cv.icon),
+        (CONF_ENTITY_CATEGORY, entity_category, cv.entity_category),
+        (CONF_DEVICE_CLASS, device_class, validate_device_class),
+    ]:
+        if default is not _UNDEF:
+            schema[cv.Optional(key, default=default)] = validator
+
+    return BUTTON_SCHEMA.extend(schema)
 
 
 async def setup_button_core_(var, config):
diff --git a/esphome/components/output/button/__init__.py b/esphome/components/output/button/__init__.py
index 4b81cedc0c..c31865ccfb 100644
--- a/esphome/components/output/button/__init__.py
+++ b/esphome/components/output/button/__init__.py
@@ -6,13 +6,16 @@ from .. import output_ns
 
 OutputButton = output_ns.class_("OutputButton", button.Button, cg.Component)
 
-CONFIG_SCHEMA = button.BUTTON_SCHEMA.extend(
-    {
-        cv.GenerateID(): cv.declare_id(OutputButton),
-        cv.Required(CONF_OUTPUT): cv.use_id(output.BinaryOutput),
-        cv.Required(CONF_DURATION): cv.positive_time_period_milliseconds,
-    }
-).extend(cv.COMPONENT_SCHEMA)
+CONFIG_SCHEMA = (
+    button.button_schema(OutputButton)
+    .extend(
+        {
+            cv.Required(CONF_OUTPUT): cv.use_id(output.BinaryOutput),
+            cv.Required(CONF_DURATION): cv.positive_time_period_milliseconds,
+        }
+    )
+    .extend(cv.COMPONENT_SCHEMA)
+)
 
 
 async def to_code(config):
diff --git a/esphome/components/template/button/__init__.py b/esphome/components/template/button/__init__.py
index 2ad5e54c80..e0101dfc8f 100644
--- a/esphome/components/template/button/__init__.py
+++ b/esphome/components/template/button/__init__.py
@@ -1,15 +1,10 @@
-import esphome.config_validation as cv
 from esphome.components import button
 
 from .. import template_ns
 
 TemplateButton = template_ns.class_("TemplateButton", button.Button)
 
-CONFIG_SCHEMA = button.BUTTON_SCHEMA.extend(
-    {
-        cv.GenerateID(): cv.declare_id(TemplateButton),
-    }
-)
+CONFIG_SCHEMA = button.button_schema(TemplateButton)
 
 
 async def to_code(config):
diff --git a/esphome/components/wake_on_lan/button.py b/esphome/components/wake_on_lan/button.py
index 2710eb3df9..778ea60cfa 100644
--- a/esphome/components/wake_on_lan/button.py
+++ b/esphome/components/wake_on_lan/button.py
@@ -12,11 +12,12 @@ WakeOnLanButton = wake_on_lan_ns.class_("WakeOnLanButton", button.Button, cg.Com
 DEPENDENCIES = ["network"]
 
 CONFIG_SCHEMA = cv.All(
-    button.BUTTON_SCHEMA.extend(cv.COMPONENT_SCHEMA).extend(
+    button.button_schema(WakeOnLanButton)
+    .extend(cv.COMPONENT_SCHEMA)
+    .extend(
         cv.Schema(
             {
                 cv.Required(CONF_TARGET_MAC_ADDRESS): cv.mac_address,
-                cv.GenerateID(): cv.declare_id(WakeOnLanButton),
             }
         ),
     ),