add SUB_BUTTON macro and ability to button schema to define the class (#4450)

* add ability to button schema to define the class

* add SUB_BUTTON macro

esphome/components/button/__init__.py

@@ -17,6 +17,7 @@ from esphome.const import (
 )
 from esphome.core import CORE, coroutine_with_priority
 from esphome.cpp_helpers import setup_entity
+from esphome.cpp_generator import MockObjClass
 
 CODEOWNERS = ["@esphome/core"]
 IS_PLATFORM_COMPONENT = True
@@ -56,11 +57,15 @@ _UNDEF = object()
 
 
 def button_schema(
+    class_: MockObjClass = _UNDEF,
+    *,
     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:

esphome/components/button/button.h

@@ -15,6 +15,13 @@ namespace button {
     } \
   }
 
+#define SUB_BUTTON(name) \
+ protected: \
+  button::Button *name##_button_{nullptr}; \
+\
+ public: \
+  void set_##name##_button(button::Button *button) { this->name##_button_ = button; }
+
 /** Base class for all buttons.
  *
  * A button is just a momentary switch that does not have a state, only a trigger.

esphome/components/copy/button/__init__.py

@@ -16,10 +16,9 @@ CopyButton = copy_ns.class_("CopyButton", button.Button, cg.Component)
 
 
 CONFIG_SCHEMA = (
-    button.button_schema()
+    button.button_schema(CopyButton)
     .extend(
         {
-            cv.GenerateID(): cv.declare_id(CopyButton),
             cv.Required(CONF_SOURCE_ID): cv.use_id(button.Button),
         }
     )

esphome/components/factory_reset/button/__init__.py

@@ -13,15 +13,12 @@ FactoryResetButton = factory_reset_ns.class_(
     "FactoryResetButton", button.Button, cg.Component
 )
 
-CONFIG_SCHEMA = (
-    button.button_schema(
-        device_class=DEVICE_CLASS_RESTART,
-        entity_category=ENTITY_CATEGORY_CONFIG,
-        icon=ICON_RESTART_ALERT,
-    )
-    .extend({cv.GenerateID(): cv.declare_id(FactoryResetButton)})
-    .extend(cv.COMPONENT_SCHEMA)
-)
+CONFIG_SCHEMA = button.button_schema(
+    FactoryResetButton,
+    device_class=DEVICE_CLASS_RESTART,
+    entity_category=ENTITY_CATEGORY_CONFIG,
+    icon=ICON_RESTART_ALERT,
+).extend(cv.COMPONENT_SCHEMA)
 
 
 async def to_code(config):

esphome/components/restart/button/__init__.py

@@ -10,13 +10,11 @@ from esphome.const import (
 restart_ns = cg.esphome_ns.namespace("restart")
 RestartButton = restart_ns.class_("RestartButton", button.Button, cg.Component)
 
-CONFIG_SCHEMA = (
-    button.button_schema(
-        device_class=DEVICE_CLASS_RESTART, entity_category=ENTITY_CATEGORY_CONFIG
-    )
-    .extend({cv.GenerateID(): cv.declare_id(RestartButton)})
-    .extend(cv.COMPONENT_SCHEMA)
-)
+CONFIG_SCHEMA = button.button_schema(
+    RestartButton,
+    device_class=DEVICE_CLASS_RESTART,
+    entity_category=ENTITY_CATEGORY_CONFIG,
+).extend(cv.COMPONENT_SCHEMA)
 
 
 async def to_code(config):

esphome/components/safe_mode/button/__init__.py

@@ -17,11 +17,11 @@ SafeModeButton = safe_mode_ns.class_("SafeModeButton", button.Button, cg.Compone
 
 CONFIG_SCHEMA = (
     button.button_schema(
+        SafeModeButton,
         device_class=DEVICE_CLASS_RESTART,
         entity_category=ENTITY_CATEGORY_CONFIG,
         icon=ICON_RESTART_ALERT,
     )
-    .extend({cv.GenerateID(): cv.declare_id(SafeModeButton)})
     .extend({cv.GenerateID(CONF_OTA): cv.use_id(OTAComponent)})
     .extend(cv.COMPONENT_SCHEMA)
 )

esphome/components/shutdown/button/__init__.py

@@ -10,11 +10,9 @@ from esphome.const import (
 shutdown_ns = cg.esphome_ns.namespace("shutdown")
 ShutdownButton = shutdown_ns.class_("ShutdownButton", button.Button, cg.Component)
 
-CONFIG_SCHEMA = (
-    button.button_schema(entity_category=ENTITY_CATEGORY_CONFIG, icon=ICON_POWER)
-    .extend({cv.GenerateID(): cv.declare_id(ShutdownButton)})
-    .extend(cv.COMPONENT_SCHEMA)
-)
+CONFIG_SCHEMA = button.button_schema(
+    ShutdownButton, entity_category=ENTITY_CATEGORY_CONFIG, icon=ICON_POWER
+).extend(cv.COMPONENT_SCHEMA)
 
 
 async def to_code(config):