Migrate Tuya fan to TuyaFanDefinition (#166344)

homeassistant/components/tuya/fan.py

@@ -4,15 +4,9 @@ from __future__ import annotations
 
 from typing import Any
 
-from tuya_device_handlers.device_wrapper.base import DeviceWrapper
-from tuya_device_handlers.device_wrapper.common import (
-    DPCodeBooleanWrapper,
-    DPCodeEnumWrapper,
-)
-from tuya_device_handlers.device_wrapper.fan import (
-    FanDirectionEnumWrapper,
-    FanSpeedEnumWrapper,
-    FanSpeedIntegerWrapper,
+from tuya_device_handlers.definition.fan import (
+    TuyaFanDefinition,
+    get_default_definition,
 )
 from tuya_device_handlers.helpers.homeassistant import TuyaFanDirection
 from tuya_sharing import CustomerDevice, Manager
@@ -28,20 +22,8 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from . import TuyaConfigEntry
-from .const import TUYA_DISCOVERY_NEW, DeviceCategory, DPCode
+from .const import TUYA_DISCOVERY_NEW, DeviceCategory
 from .entity import TuyaEntity
-from .util import get_dpcode
-
-_DIRECTION_DPCODES = (DPCode.FAN_DIRECTION,)
-_MODE_DPCODES = (DPCode.FAN_MODE, DPCode.MODE)
-_OSCILLATE_DPCODES = (DPCode.SWITCH_HORIZONTAL, DPCode.SWITCH_VERTICAL)
-_SPEED_DPCODES = (
-    DPCode.FAN_SPEED_PERCENT,
-    DPCode.FAN_SPEED,
-    DPCode.SPEED,
-    DPCode.FAN_SPEED_ENUM,
-)
-_SWITCH_DPCODES = (DPCode.SWITCH_FAN, DPCode.FAN_SWITCH, DPCode.SWITCH)
 
 TUYA_SUPPORT_TYPE: set[DeviceCategory] = {
     DeviceCategory.CS,
@@ -61,30 +43,6 @@ _HA_TO_TUYA_DIRECTION_MAPPINGS = {
 }
 
 
-def _has_a_valid_dpcode(device: CustomerDevice) -> bool:
-    """Check if the device has at least one valid DP code."""
-    properties_to_check: list[DPCode | tuple[DPCode, ...] | None] = [
-        # Main control switch
-        _SWITCH_DPCODES,
-        # Other properties
-        _SPEED_DPCODES,
-        _OSCILLATE_DPCODES,
-        _DIRECTION_DPCODES,
-    ]
-    return any(get_dpcode(device, code) for code in properties_to_check)
-
-
-def _get_speed_wrapper(
-    device: CustomerDevice,
-) -> DeviceWrapper[int] | None:
-    """Get the speed wrapper for the device."""
-    if int_wrapper := FanSpeedIntegerWrapper.find_dpcode(
-        device, _SPEED_DPCODES, prefer_function=True
-    ):
-        return int_wrapper
-    return FanSpeedEnumWrapper.find_dpcode(device, _SPEED_DPCODES, prefer_function=True)
-
-
 async def async_setup_entry(
     hass: HomeAssistant,
     entry: TuyaConfigEntry,
@@ -99,26 +57,10 @@ async def async_setup_entry(
         entities: list[TuyaFanEntity] = []
         for device_id in device_ids:
             device = manager.device_map[device_id]
-            if device.category in TUYA_SUPPORT_TYPE and _has_a_valid_dpcode(device):
-                entities.append(
-                    TuyaFanEntity(
-                        device,
-                        manager,
-                        direction_wrapper=FanDirectionEnumWrapper.find_dpcode(
-                            device, _DIRECTION_DPCODES, prefer_function=True
-                        ),
-                        mode_wrapper=DPCodeEnumWrapper.find_dpcode(
-                            device, _MODE_DPCODES, prefer_function=True
-                        ),
-                        oscillate_wrapper=DPCodeBooleanWrapper.find_dpcode(
-                            device, _OSCILLATE_DPCODES, prefer_function=True
-                        ),
-                        speed_wrapper=_get_speed_wrapper(device),
-                        switch_wrapper=DPCodeBooleanWrapper.find_dpcode(
-                            device, _SWITCH_DPCODES, prefer_function=True
-                        ),
-                    )
-                )
+            if device.category in TUYA_SUPPORT_TYPE and (
+                definition := get_default_definition(device)
+            ):
+                entities.append(TuyaFanEntity(device, manager, definition))
         async_add_entities(entities)
 
     async_discover_device([*manager.device_map])
@@ -137,38 +79,33 @@ class TuyaFanEntity(TuyaEntity, FanEntity):
         self,
         device: CustomerDevice,
         device_manager: Manager,
-        *,
-        direction_wrapper: DeviceWrapper[TuyaFanDirection] | None,
-        mode_wrapper: DeviceWrapper[str] | None,
-        oscillate_wrapper: DeviceWrapper[bool] | None,
-        speed_wrapper: DeviceWrapper[int] | None,
-        switch_wrapper: DeviceWrapper[bool] | None,
+        definition: TuyaFanDefinition,
     ) -> None:
         """Init Tuya Fan Device."""
         super().__init__(device, device_manager)
-        self._direction_wrapper = direction_wrapper
-        self._mode_wrapper = mode_wrapper
-        self._oscillate_wrapper = oscillate_wrapper
-        self._speed_wrapper = speed_wrapper
-        self._switch_wrapper = switch_wrapper
+        self._direction_wrapper = definition.direction_wrapper
+        self._mode_wrapper = definition.mode_wrapper
+        self._oscillate_wrapper = definition.oscillate_wrapper
+        self._speed_wrapper = definition.speed_wrapper
+        self._switch_wrapper = definition.switch_wrapper
 
-        if mode_wrapper:
+        if definition.mode_wrapper:
             self._attr_supported_features |= FanEntityFeature.PRESET_MODE
-            self._attr_preset_modes = mode_wrapper.options
+            self._attr_preset_modes = definition.mode_wrapper.options
 
-        if speed_wrapper:
+        if definition.speed_wrapper:
             self._attr_supported_features |= FanEntityFeature.SET_SPEED
             # if speed is from an enum, set speed count from options
             # else keep entity default 100
-            if hasattr(speed_wrapper, "options"):
-                self._attr_speed_count = len(speed_wrapper.options)
+            if hasattr(definition.speed_wrapper, "options"):
+                self._attr_speed_count = len(definition.speed_wrapper.options)
 
-        if oscillate_wrapper:
+        if definition.oscillate_wrapper:
             self._attr_supported_features |= FanEntityFeature.OSCILLATE
 
-        if direction_wrapper:
+        if definition.direction_wrapper:
             self._attr_supported_features |= FanEntityFeature.DIRECTION
-        if switch_wrapper:
+        if definition.switch_wrapper:
             self._attr_supported_features |= (
                 FanEntityFeature.TURN_ON | FanEntityFeature.TURN_OFF
             )