Change ZHA Tuya plug matches to use quirk IDs

home-assistant · Oct 25, 2023 · 40879ad · 40879ad
1 parent 7c93d4f
commit 40879ad
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 83 deletions.
diff --git a/homeassistant/components/zha/core/cluster_handlers/general.py b/homeassistant/components/zha/core/cluster_handlers/general.py
@@ -4,6 +4,7 @@
 from collections.abc import Coroutine
 from typing import TYPE_CHECKING, Any
 
+from zhaquirks.quirk_ids import TUYA_PLUG_ONOFF
 import zigpy.exceptions
 import zigpy.types as t
 import zigpy.zcl
@@ -347,26 +348,10 @@ def __init__(self, cluster: zigpy.zcl.Cluster, endpoint: Endpoint) -> None:
         super().__init__(cluster, endpoint)
         self._off_listener = None
 
-        if self.cluster.endpoint.model not in (
-            "TS011F",
-            "TS0121",
-            "TS0001",
-            "TS0002",
-            "TS0003",
-            "TS0004",
-        ):
-            return
-
-        try:
-            self.cluster.find_attribute("backlight_mode")
-        except KeyError:
-            return
-
-        self.ZCL_INIT_ATTRS = self.ZCL_INIT_ATTRS.copy()
-        self.ZCL_INIT_ATTRS["backlight_mode"] = True
-        self.ZCL_INIT_ATTRS["power_on_state"] = True
-
-        if self.cluster.endpoint.model == "TS011F":
+        if endpoint.device.quirk_id == TUYA_PLUG_ONOFF:
+            self.ZCL_INIT_ATTRS = self.ZCL_INIT_ATTRS.copy()
+            self.ZCL_INIT_ATTRS["backlight_mode"] = True
+            self.ZCL_INIT_ATTRS["power_on_state"] = True
             self.ZCL_INIT_ATTRS["child_lock"] = True
 
     @classmethod

diff --git a/homeassistant/components/zha/core/cluster_handlers/manufacturerspecific.py b/homeassistant/components/zha/core/cluster_handlers/manufacturerspecific.py
@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Any
 
 from zhaquirks.inovelli.types import AllLEDEffectType, SingleLEDEffectType
+from zhaquirks.quirk_ids import TUYA_PLUG_MANUFACTURER
 import zigpy.zcl
 
 from homeassistant.core import callback
@@ -72,25 +73,7 @@ class TuyaClusterHandler(ClusterHandler):
     def __init__(self, cluster: zigpy.zcl.Cluster, endpoint: Endpoint) -> None:
         """Initialize TuyaClusterHandler."""
         super().__init__(cluster, endpoint)
-
-        if self.cluster.endpoint.manufacturer in (
-            "_TZE200_7tdtqgwv",
-            "_TZE200_amp6tsvy",
-            "_TZE200_oisqyl4o",
-            "_TZE200_vhy3iakz",
-            "_TZ3000_uim07oem",
-            "_TZE200_wfxuhoea",
-            "_TZE200_tviaymwx",
-            "_TZE200_g1ib5ldv",
-            "_TZE200_wunufsil",
-            "_TZE200_7deq70b8",
-            "_TZE200_tz32mtza",
-            "_TZE200_2hf7x9n3",
-            "_TZE200_aqnazj70",
-            "_TZE200_1ozguk6x",
-            "_TZE200_k6jhsr0q",
-            "_TZE200_9mahtqtg",
-        ):
+        if endpoint.device.quirk_id == TUYA_PLUG_MANUFACTURER:
             self.ZCL_INIT_ATTRS = {
                 "backlight_mode": True,
                 "power_on_state": True,

diff --git a/homeassistant/components/zha/select.py b/homeassistant/components/zha/select.py
@@ -6,6 +6,7 @@
 import logging
 from typing import TYPE_CHECKING, Any, Self
 
+from zhaquirks.quirk_ids import TUYA_PLUG_MANUFACTURER, TUYA_PLUG_ONOFF
 from zigpy import types
 from zigpy.zcl.clusters.general import OnOff
 from zigpy.zcl.clusters.security import IasWd
@@ -246,29 +247,10 @@ class TuyaPowerOnState(types.enum8):
 
 
 @CONFIG_DIAGNOSTIC_MATCH(
-    cluster_handler_names=CLUSTER_HANDLER_ON_OFF,
-    models={"TS011F", "TS0121", "TS0001", "TS0002", "TS0003", "TS0004"},
+    cluster_handler_names=CLUSTER_HANDLER_ON_OFF, quirk_ids=TUYA_PLUG_ONOFF
 )
 @CONFIG_DIAGNOSTIC_MATCH(
-    cluster_handler_names="tuya_manufacturer",
-    manufacturers={
-        "_TZE200_7tdtqgwv",
-        "_TZE200_amp6tsvy",
-        "_TZE200_oisqyl4o",
-        "_TZE200_vhy3iakz",
-        "_TZ3000_uim07oem",
-        "_TZE200_wfxuhoea",
-        "_TZE200_tviaymwx",
-        "_TZE200_g1ib5ldv",
-        "_TZE200_wunufsil",
-        "_TZE200_7deq70b8",
-        "_TZE200_tz32mtza",
-        "_TZE200_2hf7x9n3",
-        "_TZE200_aqnazj70",
-        "_TZE200_1ozguk6x",
-        "_TZE200_k6jhsr0q",
-        "_TZE200_9mahtqtg",
-    },
+    cluster_handler_names="tuya_manufacturer", quirk_ids=TUYA_PLUG_MANUFACTURER
 )
 class TuyaPowerOnStateSelectEntity(ZCLEnumSelectEntity):
     """Representation of a ZHA power on state select entity."""
@@ -288,8 +270,7 @@ class TuyaBacklightMode(types.enum8):
 
 
 @CONFIG_DIAGNOSTIC_MATCH(
-    cluster_handler_names=CLUSTER_HANDLER_ON_OFF,
-    models={"TS011F", "TS0121", "TS0001", "TS0002", "TS0003", "TS0004"},
+    cluster_handler_names=CLUSTER_HANDLER_ON_OFF, quirk_ids=TUYA_PLUG_ONOFF
 )
 class TuyaBacklightModeSelectEntity(ZCLEnumSelectEntity):
     """Representation of a ZHA backlight mode select entity."""
@@ -310,25 +291,7 @@ class MoesBacklightMode(types.enum8):
 
 
 @CONFIG_DIAGNOSTIC_MATCH(
-    cluster_handler_names="tuya_manufacturer",
-    manufacturers={
-        "_TZE200_7tdtqgwv",
-        "_TZE200_amp6tsvy",
-        "_TZE200_oisqyl4o",
-        "_TZE200_vhy3iakz",
-        "_TZ3000_uim07oem",
-        "_TZE200_wfxuhoea",
-        "_TZE200_tviaymwx",
-        "_TZE200_g1ib5ldv",
-        "_TZE200_wunufsil",
-        "_TZE200_7deq70b8",
-        "_TZE200_tz32mtza",
-        "_TZE200_2hf7x9n3",
-        "_TZE200_aqnazj70",
-        "_TZE200_1ozguk6x",
-        "_TZE200_k6jhsr0q",
-        "_TZE200_9mahtqtg",
-    },
+    cluster_handler_names="tuya_manufacturer", quirk_ids=TUYA_PLUG_MANUFACTURER
 )
 class MoesBacklightModeSelectEntity(ZCLEnumSelectEntity):
     """Moes devices have a different backlight mode select options."""

diff --git a/homeassistant/components/zha/switch.py b/homeassistant/components/zha/switch.py
@@ -5,6 +5,7 @@
 import logging
 from typing import TYPE_CHECKING, Any, Self
 
+from zhaquirks.quirk_ids import TUYA_PLUG_ONOFF
 from zigpy.zcl.clusters.general import OnOff
 from zigpy.zcl.foundation import Status
 
@@ -488,8 +489,7 @@ class AqaraPetFeederChildLock(ZHASwitchConfigurationEntity):
 
 
 @CONFIG_DIAGNOSTIC_MATCH(
-    cluster_handler_names=CLUSTER_HANDLER_ON_OFF,
-    models={"TS011F"},
+    cluster_handler_names=CLUSTER_HANDLER_ON_OFF, quirk_ids=TUYA_PLUG_ONOFF
 )
 class TuyaChildLockSwitch(ZHASwitchConfigurationEntity):
     """Representation of a child lock configuration entity."""