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