Simplify button_schema function (#4468)

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):

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(
+CONFIG_SCHEMA = (
+    button.button_schema(OutputButton)
+    .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)
+    )
+    .extend(cv.COMPONENT_SCHEMA)
+)
 
 
 async def to_code(config):

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):

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),
             }
         ),
     ),