Migrate Tuya light (color_mode) to use wrapper class (home-assistant#156582)

Author: epenet

homeassistant/components/tuya/light.py

@@ -29,7 +29,12 @@
 from . import TuyaConfigEntry
 from .const import TUYA_DISCOVERY_NEW, DeviceCategory, DPCode, DPType, WorkMode
 from .entity import TuyaEntity
-from .models import DPCodeBooleanWrapper, IntegerTypeData, find_dpcode
+from .models import (
+    DPCodeBooleanWrapper,
+    DPCodeEnumWrapper,
+    IntegerTypeData,
+    find_dpcode,
+)
 from .util import get_dpcode, get_dptype, remap_value
 
 
@@ -429,7 +434,13 @@ def async_discover_device(device_ids: list[str]):
             if descriptions := LIGHTS.get(device.category):
                 entities.extend(
                     TuyaLightEntity(
-                        device, manager, description, switch_wrapper=switch_wrapper
+                        device,
+                        manager,
+                        description,
+                        color_mode_wrapper=DPCodeEnumWrapper.find_dpcode(
+                            device, description.color_mode, prefer_function=True
+                        ),
+                        switch_wrapper=switch_wrapper,
                     )
                     for description in descriptions
                     if (
@@ -458,7 +469,6 @@ class TuyaLightEntity(TuyaEntity, LightEntity):
     _brightness: IntegerTypeData | None = None
     _color_data_dpcode: DPCode | None = None
     _color_data_type: ColorTypeData | None = None
-    _color_mode: DPCode | None = None
     _color_temp: IntegerTypeData | None = None
     _white_color_mode = ColorMode.COLOR_TEMP
     _fixed_color_mode: ColorMode | None = None
@@ -471,19 +481,18 @@ def __init__(
         device_manager: Manager,
         description: TuyaLightEntityDescription,
         *,
+        color_mode_wrapper: DPCodeEnumWrapper | None,
         switch_wrapper: DPCodeBooleanWrapper,
     ) -> None:
         """Init TuyaHaLight."""
         super().__init__(device, device_manager)
         self.entity_description = description
         self._attr_unique_id = f"{super().unique_id}{description.key}"
+        self._color_mode_wrapper = color_mode_wrapper
         self._switch_wrapper = switch_wrapper
 
         color_modes: set[ColorMode] = {ColorMode.ONOFF}
 
-        # Determine DPCodes
-        self._color_mode_dpcode = get_dpcode(self.device, description.color_mode)
-
         if int_type := find_dpcode(
             self.device,
             description.brightness,
@@ -537,15 +546,8 @@ def __init__(
         # work_mode "white"
         elif (
             color_supported(color_modes)
-            and (
-                color_mode_enum := find_dpcode(
-                    self.device,
-                    description.color_mode,
-                    dptype=DPType.ENUM,
-                    prefer_function=True,
-                )
-            )
-            and WorkMode.WHITE.value in color_mode_enum.range
+            and color_mode_wrapper is not None
+            and WorkMode.WHITE in color_mode_wrapper.type_information.range
         ):
             color_modes.add(ColorMode.WHITE)
             self._white_color_mode = ColorMode.WHITE
@@ -566,14 +568,13 @@ def turn_on(self, **kwargs: Any) -> None:
             self._switch_wrapper.get_update_command(self.device, True),
         ]
 
-        if self._color_mode_dpcode and (
+        if self._color_mode_wrapper and (
             ATTR_WHITE in kwargs or ATTR_COLOR_TEMP_KELVIN in kwargs
         ):
             commands += [
-                {
-                    "code": self._color_mode_dpcode,
-                    "value": WorkMode.WHITE,
-                },
+                self._color_mode_wrapper.get_update_command(
+                    self.device, WorkMode.WHITE
+                ),
             ]
 
         if self._color_temp and ATTR_COLOR_TEMP_KELVIN in kwargs:
@@ -602,12 +603,11 @@ def turn_on(self, **kwargs: Any) -> None:
                 and ATTR_COLOR_TEMP_KELVIN not in kwargs
             )
         ):
-            if self._color_mode_dpcode:
+            if self._color_mode_wrapper:
                 commands += [
-                    {
-                        "code": self._color_mode_dpcode,
-                        "value": WorkMode.COLOUR,
-                    },
+                    self._color_mode_wrapper.get_update_command(
+                        self.device, WorkMode.COLOUR
+                    ),
                 ]
 
             if not (brightness := kwargs.get(ATTR_BRIGHTNESS)):
@@ -765,8 +765,8 @@ def color_mode(self) -> ColorMode:
         # and HS, determine which mode the light is in. We consider it to be in HS color
         # mode, when work mode is anything else than "white".
         if (
-            self._color_mode_dpcode
-            and self.device.status.get(self._color_mode_dpcode) != WorkMode.WHITE
+            self._color_mode_wrapper
+            and self._read_wrapper(self._color_mode_wrapper) != WorkMode.WHITE
         ):
             return ColorMode.HS
         return self._white_color_mode