Cleanup, set defaults

Author: prairiesnpr

zhaquirks/tuya/mcu/__init__.py

@@ -11,6 +11,7 @@
 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
 
@@ -29,6 +30,7 @@
     TuyaEnchantableCluster,
     TuyaLocalCluster,
     TuyaNewManufCluster,
+    TuyaPowerConfigurationCluster2AAA,
     TuyaTimePayload,
 )
 
@@ -730,6 +732,14 @@ class TuyaLevelControlManufCluster(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."""
+
+
 class TuyaQuirkBuilder(QuirkBuilder):
     """Tuya QuirkBuilder."""
 
@@ -741,49 +751,10 @@ def __init__(
         self.tuya_dp_to_attribute: dict[int, DPToAttributeMapping] = {}
         super().__init__(manufacturer, model, registry)
 
-    def add_tuya_dp(
+    def battery(
         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_battery_config(
-        self,
-        dp_id: int,
-        power_cfg: PowerConfiguration,
+        power_cfg: PowerConfiguration = TuyaPowerConfigurationCluster2AAA,
         converter: Optional[
             Callable[
                 [
@@ -794,7 +765,7 @@ def add_battery_config(
         ] = lambda x: x * 2,
     ) -> QuirkBuilder:
         """Add a Tuya Battery Power Configuration."""
-        self.add_tuya_dp(
+        self.tuya_dp(
             dp_id,
             power_cfg.ep_attribute,
             "battery_percentage_remaining",
@@ -803,10 +774,10 @@ def add_battery_config(
         self.adds(power_cfg)
         return self
 
-    def add_soil_moisture_config(
+    def soil_moisture(
         self,
         dp_id: int,
-        soil_cfg: TuyaLocalCluster,
+        soil_cfg: TuyaLocalCluster = TuyaSoilMoisture,
         converter: Optional[
             Callable[
                 [
@@ -826,10 +797,10 @@ def add_soil_moisture_config(
         self.adds(soil_cfg)
         return self
 
-    def add_temperature_config(
+    def temperature(
         self,
         dp_id: int,
-        temp_cfg: TuyaLocalCluster,
+        temp_cfg: TuyaLocalCluster = TuyaTemperatureMeasurement,
         converter: Optional[
             Callable[
                 [
@@ -849,6 +820,45 @@ def add_temperature_config(
         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."""
 

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, TuyaQuirkBuilder
-
-
-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):
@@ -337,12 +330,9 @@ class TuyaTempHumiditySensorVar04(CustomDevice):
 (
     TuyaQuirkBuilder("_TZE284_aao3yzhs", "TS0601")
     .also_applies_to("_TZE284_sgabhwa6", "TS0601")
-    .add_battery_config(15, TuyaPowerConfigurationCluster2AAA)
-    .add_temperature_config(5, TuyaTemperatureMeasurement, converter=lambda x: x * 10)
-    .add_battery_config(15, TuyaPowerConfigurationCluster2AAA)
-    .add_soil_moisture_config(3, TuyaSoilMoisture)
-    .adds(TuyaSoilMoisture)
-    .adds(TuyaTemperatureMeasurement)
+    .temperature(5, converter=lambda x: x * 10)
+    .battery(15)
+    .soil_moisture(3)
     .add_to_registry()
 )
 
@@ -352,8 +342,8 @@ class TuyaTempHumiditySensorVar04(CustomDevice):
     .also_applies_to("_TZE200_ga1maeof", "TS0601")
     .also_applies_to("_TZE200_9cqcpkgb", "TS0601")
     .also_applies_to("_TZE204_myd45weu", "TS0601")
-    .add_temperature_config(5, TuyaTemperatureMeasurement)
-    .add_battery_config(15, TuyaPowerConfigurationCluster2AAA)
-    .add_soil_moisture_config(3, TuyaSoilMoisture)
+    .temperature(5)
+    .battery(15)
+    .soil_moisture(3)
     .add_to_registry()
 )