Refactor

Author: prairiesnpr

zhaquirks/tuya/mcu/__init__.py

@@ -5,6 +5,7 @@
 import datetime
 from typing import Any, Optional, Union
 
+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
@@ -724,3 +725,62 @@ class TuyaLevelControlManufCluster(TuyaMCUCluster):
         17: "_dp_2_attr_update",
         18: "_dp_2_attr_update",
     }
+
+
+class TuyaQuirkBuilder(QuirkBuilder):
+    """Tuya QuirkBuilder."""
+
+    tuya_attributes = TuyaMCUCluster.attributes.copy()
+    tuya_data_point_handlers: dict[int, str] = {}
+    tuya_dp_to_attribute: dict[int, DPToAttributeMapping] = {}
+
+    def add_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",
+    ):
+        """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

@@ -1,11 +1,9 @@
 """Tuya temp and humidity sensors."""
 
-from collections.abc import Callable
-from typing import Any, Optional, Union
+from typing import Any
 
 from zigpy.profiles import zha
 from zigpy.quirks import CustomDevice
-from zigpy.quirks.v2 import QuirkBuilder, QuirksV2RegistryEntry
 from zigpy.zcl.clusters.general import Basic, Groups, Ota, Scenes, Time
 from zigpy.zcl.clusters.measurement import (
     RelativeHumidity,
@@ -23,7 +21,7 @@
     SKIP_CONFIGURATION,
 )
 from zhaquirks.tuya import TuyaLocalCluster, TuyaPowerConfigurationCluster2AAA
-from zhaquirks.tuya.mcu import DPToAttributeMapping, TuyaMCUCluster
+from zhaquirks.tuya.mcu import DPToAttributeMapping, TuyaMCUCluster, TuyaQuirkBuilder
 
 
 class TuyaTemperatureMeasurement(TemperatureMeasurement, TuyaLocalCluster):
@@ -336,65 +334,6 @@ class TuyaTempHumiditySensorVar04(CustomDevice):
     }
 
 
-class TuyaQuirkBuilder(QuirkBuilder):
-    """Tuya QuirkBuilder."""
-
-    tuya_attributes = TuyaMCUCluster.attributes.copy()
-    tuya_data_point_handlers: dict[int, str] = {}
-    tuya_dp_to_attribute: dict[int, DPToAttributeMapping] = {}
-
-    def add_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",
-    ):
-        """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()
-
-
 (
     TuyaQuirkBuilder("_TZE284_aao3yzhs", "TS0601")
     .also_applies_to("_TZE284_sgabhwa6", "TS0601")