Simplify DPCode lookup in Tuya (#150052)

homeassistant/components/tuya/alarm_control_panel.py

@@ -22,6 +22,7 @@ from . import TuyaConfigEntry
 from .const import TUYA_DISCOVERY_NEW, DPCode, DPType
 from .entity import TuyaEntity
 from .models import EnumTypeData
+from .util import get_dpcode
 
 
 @dataclass(frozen=True)
@@ -140,7 +141,7 @@ class TuyaAlarmEntity(TuyaEntity, AlarmControlPanelEntity):
             self._master_state = enum_type
 
         # Determine alarm message
-        if dp_code := self.find_dpcode(description.alarm_msg, prefer_function=True):
+        if dp_code := get_dpcode(self.device, description.alarm_msg):
             self._alarm_msg_dpcode = dp_code
 
     @property

homeassistant/components/tuya/climate.py

@@ -27,6 +27,7 @@ from . import TuyaConfigEntry
 from .const import TUYA_DISCOVERY_NEW, DPCode, DPType
 from .entity import TuyaEntity
 from .models import IntegerTypeData
+from .util import get_dpcode
 
 TUYA_HVAC_TO_HA = {
     "auto": HVACMode.HEAT_COOL,
@@ -229,7 +230,7 @@ class TuyaClimateEntity(TuyaEntity, ClimateEntity):
                 self._attr_hvac_modes.append(description.switch_only_hvac_mode)
                 self._attr_preset_modes = unknown_hvac_modes
                 self._attr_supported_features |= ClimateEntityFeature.PRESET_MODE
-        elif self.find_dpcode(DPCode.SWITCH, prefer_function=True):
+        elif get_dpcode(self.device, DPCode.SWITCH):
             self._attr_hvac_modes = [
                 HVACMode.OFF,
                 description.switch_only_hvac_mode,
@@ -261,24 +262,24 @@ class TuyaClimateEntity(TuyaEntity, ClimateEntity):
             self._fan_mode_dp_code = enum_type.dpcode
 
         # Determine swing modes
-        if self.find_dpcode(
+        if get_dpcode(
+            self.device,
             (
                 DPCode.SHAKE,
                 DPCode.SWING,
                 DPCode.SWITCH_HORIZONTAL,
                 DPCode.SWITCH_VERTICAL,
             ),
-            prefer_function=True,
         ):
             self._attr_supported_features |= ClimateEntityFeature.SWING_MODE
             self._attr_swing_modes = [SWING_OFF]
-            if self.find_dpcode((DPCode.SHAKE, DPCode.SWING), prefer_function=True):
+            if get_dpcode(self.device, (DPCode.SHAKE, DPCode.SWING)):
                 self._attr_swing_modes.append(SWING_ON)
 
-            if self.find_dpcode(DPCode.SWITCH_HORIZONTAL, prefer_function=True):
+            if get_dpcode(self.device, DPCode.SWITCH_HORIZONTAL):
                 self._attr_swing_modes.append(SWING_HORIZONTAL)
 
-            if self.find_dpcode(DPCode.SWITCH_VERTICAL, prefer_function=True):
+            if get_dpcode(self.device, DPCode.SWITCH_VERTICAL):
                 self._attr_swing_modes.append(SWING_VERTICAL)
 
         if DPCode.SWITCH in self.device.function:

homeassistant/components/tuya/cover.py

@@ -23,6 +23,7 @@ from . import TuyaConfigEntry
 from .const import TUYA_DISCOVERY_NEW, DPCode, DPType
 from .entity import TuyaEntity
 from .models import IntegerTypeData
+from .util import get_dpcode
 
 
 @dataclass(frozen=True)
@@ -202,7 +203,7 @@ class TuyaCoverEntity(TuyaEntity, CoverEntity):
         self._attr_supported_features = CoverEntityFeature(0)
 
         # Check if this cover is based on a switch or has controls
-        if self.find_dpcode(description.key, prefer_function=True):
+        if get_dpcode(self.device, description.key):
             if device.function[description.key].type == "Boolean":
                 self._attr_supported_features |= (
                     CoverEntityFeature.OPEN | CoverEntityFeature.CLOSE

homeassistant/components/tuya/entity.py

@@ -72,22 +72,17 @@ class TuyaEntity(Entity):
         dptype: Literal[DPType.INTEGER],
     ) -> IntegerTypeData | None: ...
 
-    @overload
     def find_dpcode(
         self,
         dpcodes: str | DPCode | tuple[DPCode, ...] | None,
         *,
         prefer_function: bool = False,
-    ) -> DPCode | None: ...
+        dptype: DPType,
+    ) -> EnumTypeData | IntegerTypeData | None:
+        """Find type information for a matching DP code available for this device."""
+        if dptype not in (DPType.ENUM, DPType.INTEGER):
+            raise NotImplementedError("Only ENUM and INTEGER types are supported")
 
-    def find_dpcode(
-        self,
-        dpcodes: str | DPCode | tuple[DPCode, ...] | None,
-        *,
-        prefer_function: bool = False,
-        dptype: DPType | None = None,
-    ) -> DPCode | EnumTypeData | IntegerTypeData | None:
-        """Find a matching DP code available on for this device."""
         if dpcodes is None:
             return None
 
@@ -100,11 +95,6 @@ class TuyaEntity(Entity):
         if prefer_function:
             order = ["function", "status_range"]
 
-        # When we are not looking for a specific datatype, we can append status for
-        # searching
-        if not dptype:
-            order.append("status")
-
         for dpcode in dpcodes:
             for key in order:
                 if dpcode not in getattr(self.device, key):
@@ -133,9 +123,6 @@ class TuyaEntity(Entity):
                         continue
                     return integer_type
 
-                if dptype not in (DPType.ENUM, DPType.INTEGER):
-                    return dpcode
-
         return None
 
     def get_dptype(

homeassistant/components/tuya/fan.py

@@ -24,6 +24,7 @@ from . import TuyaConfigEntry
 from .const import TUYA_DISCOVERY_NEW, DPCode, DPType
 from .entity import TuyaEntity
 from .models import EnumTypeData, IntegerTypeData
+from .util import get_dpcode
 
 TUYA_SUPPORT_TYPE = {
     # Dehumidifier
@@ -90,8 +91,8 @@ class TuyaFanEntity(TuyaEntity, FanEntity):
         """Init Tuya Fan Device."""
         super().__init__(device, device_manager)
 
-        self._switch = self.find_dpcode(
+        self._switch = get_dpcode(
-            (DPCode.SWITCH_FAN, DPCode.FAN_SWITCH, DPCode.SWITCH), prefer_function=True
+            self.device, (DPCode.SWITCH_FAN, DPCode.FAN_SWITCH, DPCode.SWITCH)
         )
 
         self._attr_preset_modes = []
@@ -120,8 +121,8 @@ class TuyaFanEntity(TuyaEntity, FanEntity):
             self._attr_supported_features |= FanEntityFeature.SET_SPEED
             self._speeds = enum_type
 
-        if dpcode := self.find_dpcode(
+        if dpcode := get_dpcode(
-            (DPCode.SWITCH_HORIZONTAL, DPCode.SWITCH_VERTICAL), prefer_function=True
+            self.device, (DPCode.SWITCH_HORIZONTAL, DPCode.SWITCH_VERTICAL)
         ):
             self._oscillate = dpcode
             self._attr_supported_features |= FanEntityFeature.OSCILLATE

homeassistant/components/tuya/humidifier.py

@@ -21,7 +21,7 @@ from . import TuyaConfigEntry
 from .const import TUYA_DISCOVERY_NEW, DPCode, DPType
 from .entity import TuyaEntity
 from .models import IntegerTypeData
-from .util import ActionDPCodeNotFoundError
+from .util import ActionDPCodeNotFoundError, get_dpcode
 
 
 @dataclass(frozen=True)
@@ -105,8 +105,8 @@ class TuyaHumidifierEntity(TuyaEntity, HumidifierEntity):
         self._attr_unique_id = f"{super().unique_id}{description.key}"
 
         # Determine main switch DPCode
-        self._switch_dpcode = self.find_dpcode(
+        self._switch_dpcode = get_dpcode(
-            description.dpcode or DPCode(description.key), prefer_function=True
+            self.device, description.dpcode or DPCode(description.key)
         )
 
         # Determine humidity parameters

homeassistant/components/tuya/light.py

@@ -29,7 +29,7 @@ from . import TuyaConfigEntry
 from .const import TUYA_DISCOVERY_NEW, DPCode, DPType, WorkMode
 from .entity import TuyaEntity
 from .models import IntegerTypeData
-from .util import remap_value
+from .util import get_dpcode, remap_value
 
 
 @dataclass
@@ -515,9 +515,7 @@ class TuyaLightEntity(TuyaEntity, LightEntity):
         color_modes: set[ColorMode] = {ColorMode.ONOFF}
 
         # Determine DPCodes
-        self._color_mode_dpcode = self.find_dpcode(
+        self._color_mode_dpcode = get_dpcode(self.device, description.color_mode)
-            description.color_mode, prefer_function=True
-        )
 
         if int_type := self.find_dpcode(
             description.brightness, dptype=DPType.INTEGER, prefer_function=True
@@ -532,7 +530,7 @@ class TuyaLightEntity(TuyaEntity, LightEntity):
             )
 
         if (
-            dpcode := self.find_dpcode(description.color_data, prefer_function=True)
+            dpcode := get_dpcode(self.device, description.color_data)
         ) and self.get_dptype(dpcode) == DPType.JSON:
             self._color_data_dpcode = dpcode
             color_modes.add(ColorMode.HS)

homeassistant/components/tuya/util.py

@@ -9,6 +9,29 @@ from homeassistant.exceptions import ServiceValidationError
 from .const import DOMAIN, DPCode
 
 
+def get_dpcode(
+    device: CustomerDevice, dpcodes: str | DPCode | tuple[DPCode, ...] | None
+) -> DPCode | None:
+    """Get the first matching DPCode from the device or return None."""
+    if dpcodes is None:
+        return None
+
+    if isinstance(dpcodes, DPCode):
+        dpcodes = (dpcodes,)
+    elif isinstance(dpcodes, str):
+        dpcodes = (DPCode(dpcodes),)
+
+    for dpcode in dpcodes:
+        if (
+            dpcode in device.function
+            or dpcode in device.status
+            or dpcode in device.status_range
+        ):
+            return dpcode
+
+    return None
+
+
 def remap_value(
     value: float,
     from_min: float = 0,

homeassistant/components/tuya/vacuum.py

@@ -19,6 +19,7 @@ from . import TuyaConfigEntry
 from .const import TUYA_DISCOVERY_NEW, DPCode, DPType
 from .entity import TuyaEntity
 from .models import EnumTypeData
+from .util import get_dpcode
 
 TUYA_MODE_RETURN_HOME = "chargego"
 TUYA_STATUS_TO_HA = {
@@ -88,11 +89,11 @@ class TuyaVacuumEntity(TuyaEntity, StateVacuumEntity):
         self._attr_supported_features = (
             VacuumEntityFeature.SEND_COMMAND | VacuumEntityFeature.STATE
         )
-        if self.find_dpcode(DPCode.PAUSE, prefer_function=True):
+        if get_dpcode(self.device, DPCode.PAUSE):
             self._attr_supported_features |= VacuumEntityFeature.PAUSE
 
         self._return_home_use_switch_charge = False
-        if self.find_dpcode(DPCode.SWITCH_CHARGE, prefer_function=True):
+        if get_dpcode(self.device, DPCode.SWITCH_CHARGE):
             self._attr_supported_features |= VacuumEntityFeature.RETURN_HOME
             self._return_home_use_switch_charge = True
         elif (
@@ -102,10 +103,10 @@ class TuyaVacuumEntity(TuyaEntity, StateVacuumEntity):
         ) and TUYA_MODE_RETURN_HOME in enum_type.range:
             self._attr_supported_features |= VacuumEntityFeature.RETURN_HOME
 
-        if self.find_dpcode(DPCode.SEEK, prefer_function=True):
+        if get_dpcode(self.device, DPCode.SEEK):
             self._attr_supported_features |= VacuumEntityFeature.LOCATE
 
-        if self.find_dpcode(DPCode.POWER_GO, prefer_function=True):
+        if get_dpcode(self.device, DPCode.POWER_GO):
             self._attr_supported_features |= (
                 VacuumEntityFeature.STOP | VacuumEntityFeature.START
             )