Reorganize CustomDevice class hierarchy

dmulcahey · Aug 14, 2024 · 9da900e · 9da900e
1 parent 284642e
commit 9da900e
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 30 deletions.
diff --git a/zigpy/quirks/__init__.py b/zigpy/quirks/__init__.py
@@ -24,6 +24,7 @@
 from zigpy.types.basic import uint16_t
 import zigpy.zcl
 from zigpy.zcl import foundation
+from zigpy.zdo import ZDO
 
 if typing.TYPE_CHECKING:
     from zigpy.application import ControllerApplication
@@ -64,17 +65,11 @@ def register_uninitialized_device_message_handler(handler: typing.Callable) -> N
         _uninitialized_device_message_handlers.append(handler)
 
 
-class CustomDevice(zigpy.device.Device):
-    """Implementation of a quirks v1 custom device."""
+class BaseCustomDevice(zigpy.device.Device):
+    """Base class for custom devices."""
 
     _copy_cluster_attr_cache = False
-
     replacement: dict[str, typing.Any] = {}
-    signature = None
-
-    def __init_subclass__(cls) -> None:
-        if getattr(cls, "signature", None) is not None:
-            _DEVICE_REGISTRY.add_to_registry(cls)
 
     def __init__(
         self,
@@ -121,6 +116,36 @@ def add_endpoint(
         self.endpoints[endpoint_id] = ep
         return ep
 
+    async def apply_custom_configuration(self, *args, **kwargs):
+        """Hook for applications to instruct instances to apply custom configuration."""
+        for endpoint in self.endpoints.values():
+            if isinstance(endpoint, ZDO):
+                continue
+            for cluster in endpoint.in_clusters.values():
+                if (
+                    isinstance(cluster, CustomCluster)
+                    and cluster.apply_custom_configuration
+                    != CustomCluster.apply_custom_configuration
+                ):
+                    await cluster.apply_custom_configuration(*args, **kwargs)
+            for cluster in endpoint.out_clusters.values():
+                if (
+                    isinstance(cluster, CustomCluster)
+                    and cluster.apply_custom_configuration
+                    != CustomCluster.apply_custom_configuration
+                ):
+                    await cluster.apply_custom_configuration(*args, **kwargs)
+
+
+class CustomDevice(BaseCustomDevice):
+    """Implementation of a quirks v1 custom device."""
+
+    signature = None
+
+    def __init_subclass__(cls) -> None:
+        if getattr(cls, "signature", None) is not None:
+            _DEVICE_REGISTRY.add_to_registry(cls)
+
 
 class CustomEndpoint(zigpy.endpoint.Endpoint):
     """Custom endpoint implementation for quirks."""

diff --git a/zigpy/quirks/v2/__init__.py b/zigpy/quirks/v2/__init__.py
@@ -21,7 +21,7 @@
     SIG_NODE_DESC,
     SIG_SKIP_CONFIG,
 )
-from zigpy.quirks import _DEVICE_REGISTRY, CustomCluster, CustomDevice, FilterType
+from zigpy.quirks import _DEVICE_REGISTRY, BaseCustomDevice, CustomCluster, FilterType
 from zigpy.quirks.registry import DeviceRegistry
 from zigpy.quirks.v2.homeassistant import EntityPlatform, EntityType
 from zigpy.quirks.v2.homeassistant.binary_sensor import BinarySensorDeviceClass
@@ -49,7 +49,7 @@
 # pylint: disable=too-few-public-methods
 
 
-class CustomDeviceV2(CustomDevice):
+class CustomDeviceV2(BaseCustomDevice):
     """Implementation of a quirks v2 custom device."""
 
     _copy_cluster_attr_cache = True
@@ -129,26 +129,6 @@ def exposes_metadata(
         """
         return self._exposes_metadata
 
-    async def apply_custom_configuration(self, *args, **kwargs):
-        """Hook for applications to instruct instances to apply custom configuration."""
-        for endpoint in self.endpoints.values():
-            if isinstance(endpoint, ZDO):
-                continue
-            for cluster in endpoint.in_clusters.values():
-                if (
-                    isinstance(cluster, CustomCluster)
-                    and cluster.apply_custom_configuration
-                    != CustomCluster.apply_custom_configuration
-                ):
-                    await cluster.apply_custom_configuration(*args, **kwargs)
-            for cluster in endpoint.out_clusters.values():
-                if (
-                    isinstance(cluster, CustomCluster)
-                    and cluster.apply_custom_configuration
-                    != CustomCluster.apply_custom_configuration
-                ):
-                    await cluster.apply_custom_configuration(*args, **kwargs)
-
 
 @attrs.define(frozen=True, kw_only=True, repr=True)
 class AddsMetadata: