From 1179a88808c2380581f71d331d6d5a8994f6f572 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Sat, 23 Apr 2022 23:04:54 +0200
Subject: [PATCH] Use ColorMode enum in yeelight (#70560)

---
 homeassistant/components/yeelight/light.py | 49 ++++++++++------------
 1 file changed, 22 insertions(+), 27 deletions(-)

diff --git a/homeassistant/components/yeelight/light.py b/homeassistant/components/yeelight/light.py
index 63302a17df5..53862a3443b 100644
--- a/homeassistant/components/yeelight/light.py
+++ b/homeassistant/components/yeelight/light.py
@@ -21,17 +21,12 @@ from homeassistant.components.light import (
     ATTR_KELVIN,
     ATTR_RGB_COLOR,
     ATTR_TRANSITION,
-    COLOR_MODE_BRIGHTNESS,
-    COLOR_MODE_COLOR_TEMP,
-    COLOR_MODE_HS,
-    COLOR_MODE_ONOFF,
-    COLOR_MODE_RGB,
-    COLOR_MODE_UNKNOWN,
     FLASH_LONG,
     FLASH_SHORT,
     SUPPORT_EFFECT,
     SUPPORT_FLASH,
     SUPPORT_TRANSITION,
+    ColorMode,
     LightEntity,
 )
 from homeassistant.config_entries import ConfigEntry
@@ -413,8 +408,8 @@ def _async_setup_services(hass: HomeAssistant):
 class YeelightGenericLight(YeelightEntity, LightEntity):
     """Representation of a Yeelight generic light."""
 
-    _attr_color_mode: str | None = COLOR_MODE_BRIGHTNESS
-    _attr_supported_color_modes: set[str] | None = {COLOR_MODE_BRIGHTNESS}
+    _attr_color_mode = ColorMode.BRIGHTNESS
+    _attr_supported_color_modes = {ColorMode.BRIGHTNESS}
     _attr_should_poll = False
 
     def __init__(self, device, entry, custom_effects=None):
@@ -640,12 +635,12 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
         if (
             not hs_color
             or not self.supported_color_modes
-            or COLOR_MODE_HS not in self.supported_color_modes
+            or ColorMode.HS not in self.supported_color_modes
         ):
             return
         if (
             not self.device.is_color_flow_enabled
-            and self.color_mode == COLOR_MODE_HS
+            and self.color_mode == ColorMode.HS
             and self.hs_color == hs_color
         ):
             _LOGGER.debug("HS already set to: %s", hs_color)
@@ -665,12 +660,12 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
         if (
             not rgb
             or not self.supported_color_modes
-            or COLOR_MODE_RGB not in self.supported_color_modes
+            or ColorMode.RGB not in self.supported_color_modes
         ):
             return
         if (
             not self.device.is_color_flow_enabled
-            and self.color_mode == COLOR_MODE_RGB
+            and self.color_mode == ColorMode.RGB
             and self.rgb_color == rgb
         ):
             _LOGGER.debug("RGB already set to: %s", rgb)
@@ -690,14 +685,14 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
         if (
             not colortemp
             or not self.supported_color_modes
-            or COLOR_MODE_COLOR_TEMP not in self.supported_color_modes
+            or ColorMode.COLOR_TEMP not in self.supported_color_modes
         ):
             return
         temp_in_k = mired_to_kelvin(colortemp)
 
         if (
             not self.device.is_color_flow_enabled
-            and self.color_mode == COLOR_MODE_COLOR_TEMP
+            and self.color_mode == ColorMode.COLOR_TEMP
             and self.color_temp == colortemp
         ):
             _LOGGER.debug("Color temp already set to: %s", temp_in_k)
@@ -882,31 +877,31 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
 class YeelightColorLightSupport(YeelightGenericLight):
     """Representation of a Color Yeelight light support."""
 
-    _attr_supported_color_modes = {COLOR_MODE_COLOR_TEMP, COLOR_MODE_HS, COLOR_MODE_RGB}
+    _attr_supported_color_modes = {ColorMode.COLOR_TEMP, ColorMode.HS, ColorMode.RGB}
 
     @property
-    def color_mode(self):
+    def color_mode(self) -> ColorMode:
         """Return the color mode."""
         color_mode = int(self._get_property("color_mode"))
         if color_mode == 1:  # RGB
-            return COLOR_MODE_RGB
+            return ColorMode.RGB
         if color_mode == 2:  # color temperature
-            return COLOR_MODE_COLOR_TEMP
+            return ColorMode.COLOR_TEMP
         if color_mode == 3:  # hsv
-            return COLOR_MODE_HS
+            return ColorMode.HS
         _LOGGER.debug("Light reported unknown color mode: %s", color_mode)
-        return COLOR_MODE_UNKNOWN
+        return ColorMode.UNKNOWN
 
     @property
     def _predefined_effects(self):
         return YEELIGHT_COLOR_EFFECT_LIST
 
 
-class YeelightWhiteTempLightSupport:
+class YeelightWhiteTempLightSupport(YeelightGenericLight):
     """Representation of a White temp Yeelight light."""
 
-    _attr_color_mode: str | None = COLOR_MODE_COLOR_TEMP
-    _attr_supported_color_modes: set[str] | None = {COLOR_MODE_COLOR_TEMP}
+    _attr_color_mode = ColorMode.COLOR_TEMP
+    _attr_supported_color_modes = {ColorMode.COLOR_TEMP}
 
     @property
     def _predefined_effects(self):
@@ -985,8 +980,8 @@ class YeelightWithNightLight(
 class YeelightNightLightMode(YeelightGenericLight):
     """Representation of a Yeelight when in nightlight mode."""
 
-    _attr_color_mode = COLOR_MODE_BRIGHTNESS
-    _attr_supported_color_modes = {COLOR_MODE_BRIGHTNESS}
+    _attr_color_mode = ColorMode.BRIGHTNESS
+    _attr_supported_color_modes = {ColorMode.BRIGHTNESS}
 
     @property
     def unique_id(self) -> str:
@@ -1037,8 +1032,8 @@ class YeelightNightLightModeWithoutBrightnessControl(YeelightNightLightMode):
     It represents case when nightlight mode brightness control is not supported.
     """
 
-    _attr_color_mode = COLOR_MODE_ONOFF
-    _attr_supported_color_modes = {COLOR_MODE_ONOFF}
+    _attr_color_mode = ColorMode.ONOFF
+    _attr_supported_color_modes = {ColorMode.ONOFF}
 
 
 class YeelightWithAmbientWithoutNightlight(YeelightWhiteTempWithoutNightlightSwitch):