Add Tuya v2 quirk builder

zhaquirks/tuya/mcu/__init__.py

@@ -5,9 +5,13 @@
 import datetime
 from typing import Any, Optional, Union
 
+from zigpy.quirks import _DEVICE_REGISTRY
+from zigpy.quirks.registry import DeviceRegistry
+from zigpy.quirks.v2 import QuirkBuilder, QuirksV2RegistryEntry
 import zigpy.types as t
 from zigpy.zcl import foundation
 from zigpy.zcl.clusters.general import LevelControl, OnOff
+from zigpy.zcl.clusters.measurement import SoilMoisture, TemperatureMeasurement
 
 from zhaquirks import Bus, DoublingPowerConfigurationCluster
 
@@ -19,12 +23,14 @@
     TUYA_SET_TIME,
     EnchantedDevice,  # noqa: F401
     NoManufacturerCluster,
+    PowerConfiguration,
     PowerOnState,
     TuyaCommand,
     TuyaDatapointData,
     TuyaEnchantableCluster,
     TuyaLocalCluster,
     TuyaNewManufCluster,
+    TuyaPowerConfigurationCluster2AAA,
     TuyaTimePayload,
 )
 
@@ -724,3 +730,143 @@ class TuyaLevelControlManufCluster(TuyaMCUCluster):
         17: "_dp_2_attr_update",
         18: "_dp_2_attr_update",
     }
+
+
+class TuyaTemperatureMeasurement(TemperatureMeasurement, TuyaLocalCluster):
+    """Tuya local TemperatureMeasurement cluster."""
+
+
+class TuyaSoilMoisture(SoilMoisture, TuyaLocalCluster):
+    """Tuya local SoilMoisture cluster with a device RH_MULTIPLIER factor if required."""
+
+
+class TuyaQuirkBuilder(QuirkBuilder):
+    """Tuya QuirkBuilder."""
+
+    def __init__(
+        self, manufacturer: str, model: str, registry: DeviceRegistry = _DEVICE_REGISTRY
+    ) -> None:
+        """Init the TuyaQuirkBuilder."""
+        self.tuya_data_point_handlers: dict[int, str] = {}
+        self.tuya_dp_to_attribute: dict[int, DPToAttributeMapping] = {}
+        super().__init__(manufacturer, model, registry)
+
+    def battery(
+        self,
+        dp_id: int = 15,
+        power_cfg: PowerConfiguration = TuyaPowerConfigurationCluster2AAA,
+        converter: Optional[
+            Callable[
+                [
+                    Any,
+                ],
+                Any,
+            ]
+        ] = lambda x: x * 2,
+    ) -> QuirkBuilder:
+        """Add a Tuya Battery Power Configuration."""
+        self.tuya_dp(
+            dp_id,
+            power_cfg.ep_attribute,
+            "battery_percentage_remaining",
+            converter=converter,
+        )
+        self.adds(power_cfg)
+        return self
+
+    def soil_moisture(
+        self,
+        dp_id: int = 3,
+        soil_cfg: TuyaLocalCluster = TuyaSoilMoisture,
+        converter: Optional[
+            Callable[
+                [
+                    Any,
+                ],
+                Any,
+            ]
+        ] = lambda x: x * 100,
+    ) -> QuirkBuilder:
+        """Add a Tuya Soil Moisture Configuration."""
+        self.tuya_dp(
+            dp_id,
+            soil_cfg.ep_attribute,
+            "measured_value",
+            converter=converter,
+        )
+        self.adds(soil_cfg)
+        return self
+
+    def temperature(
+        self,
+        dp_id: int = 5,
+        temp_cfg: TuyaLocalCluster = TuyaTemperatureMeasurement,
+        converter: Optional[
+            Callable[
+                [
+                    Any,
+                ],
+                Any,
+            ]
+        ] = lambda x: x * 100,
+    ) -> QuirkBuilder:
+        """Add a Tuya Soil Moisture Configuration."""
+        self.tuya_dp(
+            dp_id,
+            temp_cfg.ep_attribute,
+            "measured_value",
+            converter=converter,
+        )
+        self.adds(temp_cfg)
+        return self
+
+    def tuya_dp(
+        self,
+        dp_id: int,
+        ep_attribute: str,
+        attribute_name: Union[str, tuple],
+        converter: Optional[
+            Callable[
+                [
+                    Any,
+                ],
+                Any,
+            ]
+        ] = None,
+        dp_converter: Optional[
+            Callable[
+                [
+                    Any,
+                ],
+                Any,
+            ]
+        ] = None,
+        endpoint_id: Optional[int] = None,
+        dp_handler: str = "_dp_2_attr_update",
+    ) -> QuirkBuilder:
+        """Add Tuya DP Converter."""
+        self.tuya_dp_to_attribute.update(
+            {
+                dp_id: DPToAttributeMapping(
+                    ep_attribute,
+                    attribute_name,
+                    converter=converter,
+                    dp_converter=dp_converter,
+                    endpoint_id=endpoint_id,
+                )
+            }
+        )
+        self.tuya_data_point_handlers.update({dp_id: dp_handler})
+        return self
+
+    def add_to_registry(self) -> QuirksV2RegistryEntry:
+        """Build the quirks v2 registry entry."""
+
+        class TuyaReplacementCluster(TuyaMCUCluster):
+            """Replacement Tuya Cluster."""
+
+            data_point_handlers: dict[int, str] = self.tuya_data_point_handlers
+            dp_to_attribute: dict[int, DPToAttributeMapping] = self.tuya_dp_to_attribute
+
+        self.replaces(TuyaReplacementCluster)
+        return super().add_to_registry()

zhaquirks/tuya/ts0601_sensor.py

@@ -5,11 +5,7 @@
 from zigpy.profiles import zha
 from zigpy.quirks import CustomDevice
 from zigpy.zcl.clusters.general import Basic, Groups, Ota, Scenes, Time
-from zigpy.zcl.clusters.measurement import (
-    RelativeHumidity,
-    SoilMoisture,
-    TemperatureMeasurement,
-)
+from zigpy.zcl.clusters.measurement import RelativeHumidity, TemperatureMeasurement
 
 from zhaquirks.const import (
     DEVICE_TYPE,
@@ -21,15 +17,12 @@
     SKIP_CONFIGURATION,
 )
 from zhaquirks.tuya import TuyaLocalCluster, TuyaPowerConfigurationCluster2AAA
-from zhaquirks.tuya.mcu import DPToAttributeMapping, TuyaMCUCluster
-
-
-class TuyaTemperatureMeasurement(TemperatureMeasurement, TuyaLocalCluster):
-    """Tuya local TemperatureMeasurement cluster."""
-
-
-class TuyaSoilMoisture(SoilMoisture, TuyaLocalCluster):
-    """Tuya local SoilMoisture cluster with a device RH_MULTIPLIER factor if required."""
+from zhaquirks.tuya.mcu import (
+    DPToAttributeMapping,
+    TuyaMCUCluster,
+    TuyaQuirkBuilder,
+    TuyaTemperatureMeasurement,
+)
 
 
 class TuyaRelativeHumidity(RelativeHumidity, TuyaLocalCluster):
@@ -334,154 +327,23 @@ class TuyaTempHumiditySensorVar04(CustomDevice):
     }
 
 
-class SoilManufCluster(TuyaMCUCluster):
-    """Tuya Manufacturer Cluster with Temperature and Humidity data points."""
-
-    dp_to_attribute: dict[int, DPToAttributeMapping] = {
-        5: DPToAttributeMapping(
-            TuyaTemperatureMeasurement.ep_attribute,
-            "measured_value",
-            converter=lambda x: x * 100,
-        ),
-        3: DPToAttributeMapping(
-            TuyaSoilMoisture.ep_attribute,
-            "measured_value",
-            converter=lambda x: x * 100,
-        ),
-        15: DPToAttributeMapping(
-            TuyaPowerConfigurationCluster2AAA.ep_attribute,
-            "battery_percentage_remaining",
-            converter=lambda x: x * 2,  # double reported percentage
-        ),
-    }
-
-    data_point_handlers = {
-        3: "_dp_2_attr_update",
-        5: "_dp_2_attr_update",
-        15: "_dp_2_attr_update",
-    }
-
-
-class TuyaSoilSensor(CustomDevice):
-    """Tuya temp and humidity sensor (variation 03)."""
-
-    signature = {
-        # "profile_id": 260,
-        # "device_type": "0x0051",
-        # "in_clusters": ["0x0000","0x0004","0x0005","0xef00"],
-        # "out_clusters": ["0x000a","0x0019"]
-        MODELS_INFO: [
-            ("_TZE200_myd45weu", "TS0601"),
-            ("_TZE200_ga1maeof", "TS0601"),
-            ("_TZE200_9cqcpkgb", "TS0601"),
-            ("_TZE204_myd45weu", "TS0601"),
-        ],
-        ENDPOINTS: {
-            1: {
-                PROFILE_ID: zha.PROFILE_ID,
-                DEVICE_TYPE: zha.DeviceType.SMART_PLUG,
-                INPUT_CLUSTERS: [
-                    Basic.cluster_id,
-                    Groups.cluster_id,
-                    Scenes.cluster_id,
-                    TemperatureHumidityManufCluster.cluster_id,
-                ],
-                OUTPUT_CLUSTERS: [Ota.cluster_id, Time.cluster_id],
-            }
-        },
-    }
-
-    replacement = {
-        SKIP_CONFIGURATION: True,
-        ENDPOINTS: {
-            1: {
-                DEVICE_TYPE: zha.DeviceType.TEMPERATURE_SENSOR,
-                INPUT_CLUSTERS: [
-                    Basic.cluster_id,
-                    Groups.cluster_id,
-                    Scenes.cluster_id,
-                    SoilManufCluster,
-                    TuyaTemperatureMeasurement,
-                    TuyaSoilMoisture,
-                    TuyaPowerConfigurationCluster2AAA,
-                ],
-                OUTPUT_CLUSTERS: [Ota.cluster_id, Time.cluster_id],
-            }
-        },
-    }
-
-
-class SoilManufClusterVar02(TuyaMCUCluster):
-    """Tuya Manufacturer Cluster with Temperature and Humidity data points (variation 02)."""
-
-    dp_to_attribute: dict[int, DPToAttributeMapping] = {
-        5: DPToAttributeMapping(
-            TuyaTemperatureMeasurement.ep_attribute,
-            "measured_value",
-            converter=lambda x: x * 10,
-        ),
-        3: DPToAttributeMapping(
-            TuyaSoilMoisture.ep_attribute,
-            "measured_value",
-            converter=lambda x: x * 100,
-        ),
-        15: DPToAttributeMapping(
-            TuyaPowerConfigurationCluster2AAA.ep_attribute,
-            "battery_percentage_remaining",
-            converter=lambda x: x * 2,  # double reported percentage
-        ),
-    }
-
-    data_point_handlers = {
-        3: "_dp_2_attr_update",
-        5: "_dp_2_attr_update",
-        15: "_dp_2_attr_update",
-    }
-
-
-class TuyaSoilSensorVar02(CustomDevice):
-    """Tuya temp and humidity sensor (variation 05)."""
+(
+    TuyaQuirkBuilder("_TZE284_aao3yzhs", "TS0601")
+    .also_applies_to("_TZE284_sgabhwa6", "TS0601")
+    .temperature(dp_id=5, converter=lambda x: x * 10)
+    .battery(dp_id=15)
+    .soil_moisture(dp_id=3)
+    .add_to_registry()
+)
 
-    signature = {
-        # "profile_id": 260,
-        # "device_type": "0x0051",
-        # "in_clusters": ["0x0000","0x0004","0x0005","0xed00","0xef00"],
-        # "out_clusters": ["0x000a","0x0019"]
-        MODELS_INFO: [
-            ("_TZE284_aao3yzhs", "TS0601"),
-            ("_TZE284_sgabhwa6", "TS0601"),
-        ],
-        ENDPOINTS: {
-            1: {
-                PROFILE_ID: zha.PROFILE_ID,
-                DEVICE_TYPE: zha.DeviceType.SMART_PLUG,
-                INPUT_CLUSTERS: [
-                    Basic.cluster_id,
-                    Groups.cluster_id,
-                    Scenes.cluster_id,
-                    0xED00,
-                    TemperatureHumidityManufCluster.cluster_id,
-                ],
-                OUTPUT_CLUSTERS: [Ota.cluster_id, Time.cluster_id],
-            }
-        },
-    }
 
-    replacement = {
-        SKIP_CONFIGURATION: True,
-        ENDPOINTS: {
-            1: {
-                DEVICE_TYPE: zha.DeviceType.TEMPERATURE_SENSOR,
-                INPUT_CLUSTERS: [
-                    Basic.cluster_id,
-                    Groups.cluster_id,
-                    Scenes.cluster_id,
-                    SoilManufClusterVar02,
-                    TuyaTemperatureMeasurement,
-                    TuyaSoilMoisture,
-                    TuyaPowerConfigurationCluster2AAA,
-                ],
-                OUTPUT_CLUSTERS: [Ota.cluster_id, Time.cluster_id],
-            }
-        },
-    }
+(
+    TuyaQuirkBuilder("_TZE200_myd45weu", "TS0601")
+    .also_applies_to("_TZE200_ga1maeof", "TS0601")
+    .also_applies_to("_TZE200_9cqcpkgb", "TS0601")
+    .also_applies_to("_TZE204_myd45weu", "TS0601")
+    .temperature(dp_id=5)
+    .battery(dp_id=15)
+    .soil_moisture(dp_id=3)
+    .add_to_registry()
+)