Simplify DPType lookup in Tuya (home-assistant#150117)

Author: epenet

homeassistant/components/tuya/entity.py

@@ -13,16 +13,6 @@
 from .const import DOMAIN, LOGGER, TUYA_HA_SIGNAL_UPDATE_ENTITY, DPCode, DPType
 from .models import EnumTypeData, IntegerTypeData
 
-_DPTYPE_MAPPING: dict[str, DPType] = {
-    "bitmap": DPType.BITMAP,
-    "bool": DPType.BOOLEAN,
-    "enum": DPType.ENUM,
-    "json": DPType.JSON,
-    "raw": DPType.RAW,
-    "string": DPType.STRING,
-    "value": DPType.INTEGER,
-}
-
 
 class TuyaEntity(Entity):
     """Tuya base device."""
@@ -125,28 +115,6 @@ def find_dpcode(
 
         return None
 
-    def get_dptype(
-        self, dpcode: DPCode | None, *, prefer_function: bool = False
-    ) -> DPType | None:
-        """Find a matching DPCode data type available on for this device."""
-        if dpcode is None:
-            return None
-
-        order = ["status_range", "function"]
-        if prefer_function:
-            order = ["function", "status_range"]
-        for key in order:
-            if dpcode in getattr(self.device, key):
-                current_type = getattr(self.device, key)[dpcode].type
-                try:
-                    return DPType(current_type)
-                except ValueError:
-                    # Sometimes, we get ill-formed DPTypes from the cloud,
-                    # this fixes them and maps them to the correct DPType.
-                    return _DPTYPE_MAPPING.get(current_type)
-
-        return None
-
     async def async_added_to_hass(self) -> None:
         """Call when entity is added to hass."""
         self.async_on_remove(

homeassistant/components/tuya/light.py

@@ -29,7 +29,7 @@
 from .const import TUYA_DISCOVERY_NEW, DeviceCategory, DPCode, DPType, WorkMode
 from .entity import TuyaEntity
 from .models import IntegerTypeData
-from .util import get_dpcode, remap_value
+from .util import get_dpcode, get_dptype, remap_value
 
 
 @dataclass
@@ -478,9 +478,9 @@ def __init__(
                 description.brightness_min, dptype=DPType.INTEGER
             )
 
-        if (
-            dpcode := get_dpcode(self.device, description.color_data)
-        ) and self.get_dptype(dpcode, prefer_function=True) == DPType.JSON:
+        if (dpcode := get_dpcode(self.device, description.color_data)) and (
+            get_dptype(self.device, dpcode, prefer_function=True) == DPType.JSON
+        ):
             self._color_data_dpcode = dpcode
             color_modes.add(ColorMode.HS)
             if dpcode in self.device.function:

homeassistant/components/tuya/sensor.py

@@ -44,6 +44,7 @@
 )
 from .entity import TuyaEntity
 from .models import ComplexValue, ElectricityValue, EnumTypeData, IntegerTypeData
+from .util import get_dptype
 
 _WIND_DIRECTIONS = {
     "north": 0.0,
@@ -1689,7 +1690,7 @@ def __init__(
             self._type_data = enum_type
             self._type = DPType.ENUM
         else:
-            self._type = self.get_dptype(DPCode(description.key))
+            self._type = get_dptype(self.device, DPCode(description.key))
 
         # Logic to ensure the set device class and API received Unit Of Measurement
         # match Home Assistants requirements.

homeassistant/components/tuya/util.py

@@ -6,7 +6,17 @@
 
 from homeassistant.exceptions import ServiceValidationError
 
-from .const import DOMAIN, DPCode
+from .const import DOMAIN, DPCode, DPType
+
+_DPTYPE_MAPPING: dict[str, DPType] = {
+    "bitmap": DPType.BITMAP,
+    "bool": DPType.BOOLEAN,
+    "enum": DPType.ENUM,
+    "json": DPType.JSON,
+    "raw": DPType.RAW,
+    "string": DPType.STRING,
+    "value": DPType.INTEGER,
+}
 
 
 def get_dpcode(
@@ -32,6 +42,32 @@ def get_dpcode(
     return None
 
 
+def get_dptype(
+    device: CustomerDevice, dpcode: DPCode | None, *, prefer_function: bool = False
+) -> DPType | None:
+    """Find a matching DPType type information for this device DPCode."""
+    if dpcode is None:
+        return None
+
+    lookup_tuple = (
+        (device.function, device.status_range)
+        if prefer_function
+        else (device.status_range, device.function)
+    )
+
+    for device_specs in lookup_tuple:
+        if current_definition := device_specs.get(dpcode):
+            current_type = current_definition.type
+            try:
+                return DPType(current_type)
+            except ValueError:
+                # Sometimes, we get ill-formed DPTypes from the cloud,
+                # this fixes them and maps them to the correct DPType.
+                return _DPTYPE_MAPPING.get(current_type)
+
+    return None
+
+
 def remap_value(
     value: float,
     from_min: float = 0,