Use EntityDescription in Tuya Switch platform (#57581)

homeassistant/components/tuya/const.py

@@ -64,6 +64,7 @@ class DPCode(str, Enum):
     ANION = "anion"  # Ionizer unit
     BRIGHT_VALUE = "bright_value"  # Brightness
     C_F = "c_f"  # Temperature unit switching
+    CHILD_LOCK = "child_lock"  # Child lock
     COLOUR_DATA = "colour_data"  # Colored light mode
     COLOUR_DATA_V2 = "colour_data_v2"  # Colored light mode
     FAN_DIRECTION = "fan_direction"  # Fan direction
@@ -81,9 +82,24 @@ class DPCode(str, Enum):
     START = "start"  # Start
     SWING = "swing"  # Swing mode
     SWITCH = "switch"  # Switch
+    SWITCH_1 = "switch_1"  # Switch 1
+    SWITCH_2 = "switch_2"  # Switch 2
+    SWITCH_3 = "switch_3"  # Switch 3
+    SWITCH_4 = "switch_4"  # Switch 4
+    SWITCH_5 = "switch_5"  # Switch 5
+    SWITCH_6 = "switch_6"  # Switch 6
+    SWITCH_BACKLIGHT = "switch_backlight"  # Backlight switch
     SWITCH_HORIZONTAL = "switch_horizontal"  # Horizontal swing flap switch
     SWITCH_LED = "switch_led"  # Switch
+    SWITCH_SPRAY = "switch_spray"  # Spraying switch
+    SWITCH_USB1 = "switch_usb1"  # USB 1
+    SWITCH_USB2 = "switch_usb2"  # USB 2
+    SWITCH_USB3 = "switch_usb3"  # USB 3
+    SWITCH_USB4 = "switch_usb4"  # USB 4
+    SWITCH_USB5 = "switch_usb5"  # USB 5
+    SWITCH_USB6 = "switch_usb6"  # USB 6
     SWITCH_VERTICAL = "switch_vertical"  # Vertical swing flap switch
+    SWITCH_VOICE = "switch_voice"  # Voice switch
     TEMP_CURRENT = "temp_current"  # Current temperature in °C
     TEMP_CURRENT_F = "temp_current_f"  # Current temperature in °F
     TEMP_SET = "temp_set"  # Set the temperature in °C
@@ -91,6 +107,7 @@ class DPCode(str, Enum):
     TEMP_UNIT_CONVERT = "temp_unit_convert"  # Temperature unit switching
     TEMP_VALUE = "temp_value"  # Color temperature
     UV = "uv"  # UV sterilization
+    WARM = "warm"  # Heat preservation
     WATER_RESET = "water_reset"  # Resetting of water usage days
     WET = "wet"  # Humidification
     WORK_MODE = "work_mode"  # Working mode

homeassistant/components/tuya/switch.py

@@ -5,7 +5,11 @@ from typing import Any
 
 from tuya_iot import TuyaDevice, TuyaDeviceManager
 
-from homeassistant.components.switch import SwitchEntity
+from homeassistant.components.switch import (
+    DEVICE_CLASS_OUTLET,
+    SwitchEntity,
+    SwitchEntityDescription,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -15,18 +19,248 @@ from . import HomeAssistantTuyaData
 from .base import TuyaHaEntity
 from .const import DOMAIN, TUYA_DISCOVERY_NEW, DPCode
 
+# All descriptions can be found here. Mostly the Boolean data types in the
+# default instruction set of each category end up being a Switch.
 # https://developer.tuya.com/en/docs/iot/standarddescription?id=K9i5ql6waswzq
-TUYA_SUPPORT_TYPE = {
-    "kg",  # Switch
-    "cz",  # Socket
-    "pc",  # Power Strip
-    "bh",  # Smart Kettle
-    "dlq",  # Breaker
-    "cwysj",  # Pet Water Feeder
-    "kj",  # Air Purifier
-    "xxj",  # Diffuser
+SWITCHES: dict[str, tuple[SwitchEntityDescription, ...]] = {
+    # Smart Kettle
+    # https://developer.tuya.com/en/docs/iot/fbh?id=K9gf484m21yq7
+    "bh": (
+        SwitchEntityDescription(
+            key=DPCode.START,
+            name="Start",
+            icon="mdi:kettle-steam",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.WARM,
+            name="Heat preservation",
+        ),
+    ),
+    # Pet Water Feeder
+    # https://developer.tuya.com/en/docs/iot/f?id=K9gf46aewxem5
+    "cwysj": (
+        SwitchEntityDescription(
+            key=DPCode.FILTER_RESET,
+            name="Filter reset",
+            icon="mdi:filter",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.PUMP_RESET,
+            name="Water pump reset",
+            icon="mdi:pump",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH,
+            name="Power",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.WATER_RESET,
+            name="Reset of water usage days",
+            icon="mdi:water-sync",
+            entity_registry_enabled_default=False,
+        ),
+    ),
+    # Cirquit Breaker
+    "dlq": (
+        SwitchEntityDescription(
+            key=DPCode.CHILD_LOCK,
+            name="Child Lock",
+            icon="mdi:account-lock",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_1,
+            name="Switch",
+        ),
+    ),
+    # Switch
+    # https://developer.tuya.com/en/docs/iot/s?id=K9gf7o5prgf7s
+    "kg": (
+        SwitchEntityDescription(
+            key=DPCode.CHILD_LOCK,
+            name="Child Lock",
+            icon="mdi:account-lock",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_1,
+            name="Switch 1",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_2,
+            name="Switch 2",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_3,
+            name="Switch 3",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_4,
+            name="Switch 4",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_5,
+            name="Switch 5",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_6,
+            name="Switch 6",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB1,
+            name="USB 1",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB2,
+            name="USB 2",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB3,
+            name="USB 3",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB4,
+            name="USB 4",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB5,
+            name="USB 5",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB6,
+            name="USB 6",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH,
+            name="Switch",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+    ),
+    # Air Purifier
+    # https://developer.tuya.com/en/docs/iot/f?id=K9gf46h2s6dzm
+    "kj": (
+        SwitchEntityDescription(
+            key=DPCode.ANION,
+            name="Ionizer",
+            icon="mdi:minus-circle-outline",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.FILTER_RESET,
+            name="Filter cartridge reset",
+            icon="mdi:filter",
+            entity_registry_enabled_default=False,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.LOCK,
+            name="Child lock",
+            icon="mdi:account-lock",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH,
+            name="Power",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.WET,
+            name="Humidification",
+            icon="mdi:water-percent",
+        ),
+    ),
+    # Power Socket
+    # https://developer.tuya.com/en/docs/iot/s?id=K9gf7o5prgf7s
+    "pc": (
+        SwitchEntityDescription(
+            key=DPCode.CHILD_LOCK,
+            name="Child Lock",
+            icon="mdi:account-lock",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_1,
+            name="Socket 1",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_2,
+            name="Socket 2",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_3,
+            name="Socket 3",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_4,
+            name="Socket 4",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_5,
+            name="Socket 5",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_6,
+            name="Socket 6",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB1,
+            name="USB 1",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB2,
+            name="USB 2",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB3,
+            name="USB 3",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB4,
+            name="USB 4",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB5,
+            name="USB 5",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_USB6,
+            name="USB 6",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH,
+            name="Socket",
+            device_class=DEVICE_CLASS_OUTLET,
+        ),
+    ),
+    # Diffuser
+    "xxj": (
+        SwitchEntityDescription(
+            key=DPCode.SWITCH,
+            name="Power",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_SPRAY,
+            name="Spray",
+            icon="mdi:spray",
+        ),
+        SwitchEntityDescription(
+            key=DPCode.SWITCH_VOICE,
+            name="Voice",
+            icon="mdi:account-voice",
+            entity_registry_enabled_default=False,
+        ),
+    ),
 }
 
+# Socket (duplicate of `pc`)
+# https://developer.tuya.com/en/docs/iot/s?id=K9gf7o5prgf7s
+SWITCHES["cz"] = SWITCHES["pc"]
+
 
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
@@ -37,9 +271,20 @@ async def async_setup_entry(
     @callback
     def async_discover_device(device_ids: list[str]) -> None:
         """Discover and add a discovered tuya sensor."""
-        async_add_entities(
-            _setup_entities(hass, entry, hass_data.device_manager, device_ids)
-        )
+        entities: list[TuyaHaSwitch] = []
+        for device_id in device_ids:
+            device = hass_data.device_manager.device_map[device_id]
+            if descriptions := SWITCHES.get(device.category):
+                for description in descriptions:
+                    if (
+                        description.key in device.function
+                        or description.key in device.status
+                    ):
+                        entities.append(
+                            TuyaHaSwitch(device, hass_data.device_manager, description)
+                        )
+
+        async_add_entities(entities)
 
     async_discover_device([*hass_data.device_manager.device_map])
 
@@ -48,86 +293,34 @@ async def async_setup_entry(
     )
 
 
-def _setup_entities(
-    hass: HomeAssistant,
-    entry: ConfigEntry,
-    device_manager: TuyaDeviceManager,
-    device_ids: list[str],
-) -> list[TuyaHaSwitch]:
-    """Set up Tuya Switch device."""
-    entities: list[TuyaHaSwitch] = []
-    for device_id in device_ids:
-        device = device_manager.device_map[device_id]
-        if device is None or device.category not in TUYA_SUPPORT_TYPE:
-            continue
-
-        for function in device.function:
-            if device.category == "kj":
-                if function in [
-                    DPCode.ANION,
-                    DPCode.FILTER_RESET,
-                    DPCode.LIGHT,
-                    DPCode.LOCK,
-                    DPCode.UV,
-                    DPCode.WET,
-                ]:
-                    entities.append(TuyaHaSwitch(device, device_manager, function))
-
-            elif device.category == "cwysj":
-                if (
-                    function
-                    in [
-                        DPCode.FILTER_RESET,
-                        DPCode.UV,
-                        DPCode.PUMP_RESET,
-                        DPCode.WATER_RESET,
-                    ]
-                    or function.startswith(DPCode.SWITCH)
-                ):
-                    entities.append(TuyaHaSwitch(device, device_manager, function))
-
-            elif function.startswith(DPCode.START) or function.startswith(
-                DPCode.SWITCH
-            ):
-                entities.append(TuyaHaSwitch(device, device_manager, function))
-
-    return entities
-
-
 class TuyaHaSwitch(TuyaHaEntity, SwitchEntity):
     """Tuya Switch Device."""
 
-    dp_code_switch = DPCode.SWITCH
-    dp_code_start = DPCode.START
-
     def __init__(
-        self, device: TuyaDevice, device_manager: TuyaDeviceManager, dp_code: str = ""
+        self,
+        device: TuyaDevice,
+        device_manager: TuyaDeviceManager,
+        description: SwitchEntityDescription,
     ) -> None:
         """Init TuyaHaSwitch."""
         super().__init__(device, device_manager)
-
-        self.dp_code = dp_code
-        self.channel = (
-            dp_code.replace(DPCode.SWITCH, "")
-            if dp_code.startswith(DPCode.SWITCH)
-            else dp_code
-        )
-        self._attr_unique_id = f"{super().unique_id}{self.channel}"
+        self.entity_description = description
+        self._attr_unique_id = f"{super().unique_id}{description.key}"
 
     @property
     def name(self) -> str | None:
         """Return Tuya device name."""
-        return f"{self.tuya_device.name}{self.channel}"
+        return f"{self.tuya_device.name} {self.entity_description.name}"
 
     @property
     def is_on(self) -> bool:
         """Return true if switch is on."""
-        return self.tuya_device.status.get(self.dp_code, False)
+        return self.tuya_device.status.get(self.entity_description.key, False)
 
     def turn_on(self, **kwargs: Any) -> None:
         """Turn the switch on."""
-        self._send_command([{"code": self.dp_code, "value": True}])
+        self._send_command([{"code": self.entity_description.key, "value": True}])
 
     def turn_off(self, **kwargs: Any) -> None:
         """Turn the switch off."""
-        self._send_command([{"code": self.dp_code, "value": False}])
+        self._send_command([{"code": self.entity_description.key, "value": False}])