[switch] Fix schema generation (#8774)

esphome/components/dfrobot_sen0395/switch/__init__.py

@@ -2,6 +2,7 @@ import esphome.codegen as cg
 from esphome.components import switch
 import esphome.config_validation as cv
 from esphome.const import CONF_TYPE, ENTITY_CATEGORY_CONFIG
+from esphome.cpp_generator import MockObjClass
 
 from .. import CONF_DFROBOT_SEN0395_ID, DfrobotSen0395Component
 
@@ -26,32 +27,30 @@ Sen0395StartAfterBootSwitch = dfrobot_sen0395_ns.class_(
     "Sen0395StartAfterBootSwitch", DfrobotSen0395Switch
 )
 
-_SWITCH_SCHEMA = (
+
+def _switch_schema(class_: MockObjClass) -> cv.Schema:
+    return (
         switch.switch_schema(
+            class_,
             entity_category=ENTITY_CATEGORY_CONFIG,
         )
         .extend(
             {
-            cv.GenerateID(CONF_DFROBOT_SEN0395_ID): cv.use_id(DfrobotSen0395Component),
+                cv.GenerateID(CONF_DFROBOT_SEN0395_ID): cv.use_id(
+                    DfrobotSen0395Component
+                ),
             }
         )
         .extend(cv.COMPONENT_SCHEMA)
-)
+    )
+
 
 CONFIG_SCHEMA = cv.typed_schema(
     {
-        "sensor_active": _SWITCH_SCHEMA.extend(
-            {cv.GenerateID(): cv.declare_id(Sen0395PowerSwitch)}
-        ),
-        "turn_on_led": _SWITCH_SCHEMA.extend(
-            {cv.GenerateID(): cv.declare_id(Sen0395LedSwitch)}
-        ),
-        "presence_via_uart": _SWITCH_SCHEMA.extend(
-            {cv.GenerateID(): cv.declare_id(Sen0395UartPresenceSwitch)}
-        ),
-        "start_after_boot": _SWITCH_SCHEMA.extend(
-            {cv.GenerateID(): cv.declare_id(Sen0395StartAfterBootSwitch)}
-        ),
+        "sensor_active": _switch_schema(Sen0395PowerSwitch),
+        "turn_on_led": _switch_schema(Sen0395LedSwitch),
+        "presence_via_uart": _switch_schema(Sen0395UartPresenceSwitch),
+        "start_after_boot": _switch_schema(Sen0395StartAfterBootSwitch),
     }
 )
 

esphome/components/switch/__init__.py

@@ -72,6 +72,9 @@ _SWITCH_SCHEMA = (
         {
             cv.OnlyWith(CONF_MQTT_ID, "mqtt"): cv.declare_id(mqtt.MQTTSwitchComponent),
             cv.Optional(CONF_INVERTED): cv.boolean,
+            cv.Optional(CONF_RESTORE_MODE, default="ALWAYS_OFF"): cv.enum(
+                RESTORE_MODES, upper=True, space="_"
+            ),
             cv.Optional(CONF_ON_TURN_ON): automation.validate_automation(
                 {
                     cv.GenerateID(CONF_TRIGGER_ID): cv.declare_id(SwitchTurnOnTrigger),
@@ -89,54 +92,41 @@ _SWITCH_SCHEMA = (
 
 
 def switch_schema(
-    class_: MockObjClass = cv.UNDEFINED,
+    class_: MockObjClass,
     *,
-    entity_category: str = cv.UNDEFINED,
-    device_class: str = cv.UNDEFINED,
-    icon: str = cv.UNDEFINED,
     block_inverted: bool = False,
-    default_restore_mode: str = "ALWAYS_OFF",
+    default_restore_mode: str = cv.UNDEFINED,
+    device_class: str = cv.UNDEFINED,
+    entity_category: str = cv.UNDEFINED,
+    icon: str = cv.UNDEFINED,
 ):
-    schema = _SWITCH_SCHEMA.extend(
-        {
-            cv.Optional(CONF_RESTORE_MODE, default=default_restore_mode): cv.enum(
-                RESTORE_MODES, upper=True, space="_"
+    schema = {cv.GenerateID(): cv.declare_id(class_)}
+
+    for key, default, validator in [
+        (CONF_DEVICE_CLASS, device_class, validate_device_class),
+        (CONF_ENTITY_CATEGORY, entity_category, cv.entity_category),
+        (CONF_ICON, icon, cv.icon),
+        (
+            CONF_RESTORE_MODE,
+            default_restore_mode,
+            cv.enum(RESTORE_MODES, upper=True, space="_")
+            if default_restore_mode is not cv.UNDEFINED
+            else cv.UNDEFINED,
         ),
-        }
-    )
-    if class_ is not cv.UNDEFINED:
-        schema = schema.extend({cv.GenerateID(): cv.declare_id(class_)})
-    if entity_category is not cv.UNDEFINED:
-        schema = schema.extend(
-            {
-                cv.Optional(
-                    CONF_ENTITY_CATEGORY, default=entity_category
-                ): cv.entity_category
-            }
-        )
-    if device_class is not cv.UNDEFINED:
-        schema = schema.extend(
-            {
-                cv.Optional(
-                    CONF_DEVICE_CLASS, default=device_class
-                ): validate_device_class
-            }
-        )
-    if icon is not cv.UNDEFINED:
-        schema = schema.extend({cv.Optional(CONF_ICON, default=icon): cv.icon})
+    ]:
+        if default is not cv.UNDEFINED:
+            schema[cv.Optional(key, default=default)] = validator
+
     if block_inverted:
-        schema = schema.extend(
-            {
-                cv.Optional(CONF_INVERTED): cv.invalid(
+        schema[cv.Optional(CONF_INVERTED)] = cv.invalid(
             "Inverted is not supported for this platform!"
         )
-            }
-        )
-    return schema
+
+    return _SWITCH_SCHEMA.extend(schema)
 
 
 # Remove before 2025.11.0
-SWITCH_SCHEMA = switch_schema()
+SWITCH_SCHEMA = switch_schema(Switch)
 SWITCH_SCHEMA.add_extra(cv.deprecated_schema_constant("switch"))
 
 

tests/components/dfrobot_sen0395/common.yaml

@@ -26,3 +26,17 @@ dfrobot_sen0395:
 binary_sensor:
   - platform: dfrobot_sen0395
     id: mmwave_detected
+
+switch:
+  - platform: dfrobot_sen0395
+    type: sensor_active
+    id: mmwave_sensor_active
+  - platform: dfrobot_sen0395
+    type: turn_on_led
+    id: mmwave_turn_on_led
+  - platform: dfrobot_sen0395
+    type: presence_via_uart
+    id: mmwave_presence_via_uart
+  - platform: dfrobot_sen0395
+    type: start_after_boot
+    id: mmwave_start_after_boot