Use dpcode_wrapper in tuya sensor platform (home-assistant#156277)

Author: epenet

homeassistant/components/tuya/models.py

@@ -6,13 +6,14 @@
 import base64
 from dataclasses import dataclass
 import json
-import struct
 from typing import Any, Literal, Self, overload
 
 from tuya_sharing import CustomerDevice
 
+from homeassistant.util.json import json_loads
+
 from .const import DPCode, DPType
-from .util import remap_value
+from .util import parse_dptype, remap_value
 
 
 @dataclass
@@ -134,6 +135,8 @@ def from_json(cls, dpcode: DPCode, data: str) -> Self | None:
     DPType.BOOLEAN: TypeInformation,
     DPType.ENUM: EnumTypeData,
     DPType.INTEGER: IntegerTypeData,
+    DPType.JSON: TypeInformation,
+    DPType.RAW: TypeInformation,
 }
 
 
@@ -144,6 +147,9 @@ class DPCodeWrapper(ABC):
     access read conversion routines.
     """
 
+    native_unit: str | None = None
+    suggested_unit: str | None = None
+
     def __init__(self, dpcode: str) -> None:
         """Init DPCodeWrapper."""
         self.dpcode = dpcode
@@ -210,6 +216,20 @@ def find_dpcode(
         return None
 
 
+class DPCodeBase64Wrapper(DPCodeTypeInformationWrapper[TypeInformation]):
+    """Wrapper to extract information from a RAW/binary value."""
+
+    DPTYPE = DPType.RAW
+
+    def read_bytes(self, device: CustomerDevice) -> bytes | None:
+        """Read the device value for the dpcode."""
+        if (raw_value := self._read_device_status_raw(device)) is None or (
+            len(decoded := base64.b64decode(raw_value)) == 0
+        ):
+            return None
+        return decoded
+
+
 class DPCodeBooleanWrapper(DPCodeTypeInformationWrapper[TypeInformation]):
     """Simple wrapper for boolean values.
 
@@ -235,6 +255,18 @@ def _convert_value_to_raw_value(
         raise ValueError(f"Invalid boolean value `{value}`")
 
 
+class DPCodeJsonWrapper(DPCodeTypeInformationWrapper[TypeInformation]):
+    """Wrapper to extract information from a JSON value."""
+
+    DPTYPE = DPType.JSON
+
+    def read_json(self, device: CustomerDevice) -> Any | None:
+        """Read the device value for the dpcode."""
+        if (raw_value := self._read_device_status_raw(device)) is None:
+            return None
+        return json_loads(raw_value)
+
+
 class DPCodeEnumWrapper(DPCodeTypeInformationWrapper[EnumTypeData]):
     """Simple wrapper for EnumTypeData values."""
 
@@ -268,6 +300,11 @@ class DPCodeIntegerWrapper(DPCodeTypeInformationWrapper[IntegerTypeData]):
 
     DPTYPE = DPType.INTEGER
 
+    def __init__(self, dpcode: str, type_information: IntegerTypeData) -> None:
+        """Init DPCodeIntegerWrapper."""
+        super().__init__(dpcode, type_information)
+        self.native_unit = type_information.unit
+
     def read_device_status(self, device: CustomerDevice) -> float | None:
         """Read the device value for the dpcode.
 
@@ -352,6 +389,16 @@ def find_dpcode(
 ) -> IntegerTypeData | None: ...
 
 
+@overload
+def find_dpcode(
+    device: CustomerDevice,
+    dpcodes: str | DPCode | tuple[DPCode, ...] | None,
+    *,
+    prefer_function: bool = False,
+    dptype: Literal[DPType.BOOLEAN, DPType.JSON, DPType.RAW],
+) -> TypeInformation | None: ...
+
+
 def find_dpcode(
     device: CustomerDevice,
     dpcodes: str | DPCode | tuple[DPCode, ...] | None,
@@ -381,7 +428,7 @@ def find_dpcode(
         for device_specs in lookup_tuple:
             if (
                 (current_definition := device_specs.get(dpcode))
-                and current_definition.type == dptype
+                and parse_dptype(current_definition.type) is dptype
                 and (
                     type_information := type_information_cls.from_json(
                         dpcode, current_definition.values
@@ -391,44 +438,3 @@ def find_dpcode(
                 return type_information
 
     return None
-
-
-class ComplexValue:
-    """Complex value (for JSON/RAW parsing)."""
-
-    @classmethod
-    def from_json(cls, data: str) -> Self:
-        """Load JSON string and return a ComplexValue object."""
-        raise NotImplementedError("from_json is not implemented for this type")
-
-    @classmethod
-    def from_raw(cls, data: str) -> Self | None:
-        """Decode base64 string and return a ComplexValue object."""
-        raise NotImplementedError("from_raw is not implemented for this type")
-
-
-@dataclass
-class ElectricityValue(ComplexValue):
-    """Electricity complex value."""
-
-    electriccurrent: str | None = None
-    power: str | None = None
-    voltage: str | None = None
-
-    @classmethod
-    def from_json(cls, data: str) -> Self:
-        """Load JSON string and return a ElectricityValue object."""
-        return cls(**json.loads(data.lower()))
-
-    @classmethod
-    def from_raw(cls, data: str) -> Self | None:
-        """Decode base64 string and return a ElectricityValue object."""
-        raw = base64.b64decode(data)
-        if len(raw) == 0:
-            return None
-        voltage = struct.unpack(">H", raw[0:2])[0] / 10.0
-        electriccurrent = struct.unpack(">L", b"\x00" + raw[2:5])[0] / 1000.0
-        power = struct.unpack(">L", b"\x00" + raw[5:8])[0] / 1000.0
-        return cls(
-            electriccurrent=str(electriccurrent), power=str(power), voltage=str(voltage)
-        )

homeassistant/components/tuya/number.py

@@ -502,14 +502,19 @@ def __init__(
         self._attr_native_min_value = dpcode_wrapper.type_information.min_scaled
         self._attr_native_step = dpcode_wrapper.type_information.step_scaled
         if description.native_unit_of_measurement is None:
-            self._attr_native_unit_of_measurement = dpcode_wrapper.type_information.unit
+            self._attr_native_unit_of_measurement = dpcode_wrapper.native_unit
+
+        self._validate_device_class_unit()
+
+    def _validate_device_class_unit(self) -> None:
+        """Validate device class unit compatibility."""
 
         # Logic to ensure the set device class and API received Unit Of Measurement
         # match Home Assistants requirements.
         if (
             self.device_class is not None
             and not self.device_class.startswith(DOMAIN)
-            and description.native_unit_of_measurement is None
+            and self.entity_description.native_unit_of_measurement is None
             # we do not need to check mappings if the API UOM is allowed
             and self.native_unit_of_measurement
             not in NUMBER_DEVICE_CLASS_UNITS[self.device_class]