From 16a999f35f1b63059f043fe69b13d7e284027543 Mon Sep 17 00:00:00 2001
From: David Mulcahey <david.mulcahey@me.com>
Date: Sun, 11 Aug 2024 15:23:42 -0400
Subject: [PATCH] Add new quirks v2 metadata type that allows replacing all
 instances of a cluster on a device

---
 zigpy/quirks/v2/__init__.py | 69 +++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/zigpy/quirks/v2/__init__.py b/zigpy/quirks/v2/__init__.py
index 126d5f508..0a0f8d397 100644
--- a/zigpy/quirks/v2/__init__.py
+++ b/zigpy/quirks/v2/__init__.py
@@ -83,6 +83,11 @@ def __init__(
         for replace_meta in quirk_metadata.replaces_metadata:
             replace_meta(self)
 
+        for (
+            replace_occurrences_meta
+        ) in quirk_metadata.replaces_cluster_occurrences_metadata:
+            replace_occurrences_meta(self)
+
         for entity_meta in quirk_metadata.entity_metadata:
             entity_meta(self)
 
@@ -215,6 +220,30 @@ def __call__(self, device: CustomDeviceV2) -> None:
         self.add(device)
 
 
+@attrs.define(frozen=True, kw_only=True, repr=True)
+class ReplaceClusterOccurrencesMetadata:
+    """Replaces metadata for replacing all occurrences of a cluster on a device."""
+
+    cluster_types: list[ClusterType] = attrs.field()
+    cluster: type[Cluster | CustomCluster] = attrs.field()
+
+    def __call__(self, device: CustomDeviceV2) -> None:
+        """Process the replace."""
+        for endpoint in device.endpoints.values():
+            if (
+                ClusterType.Server in self.cluster_types
+                and self.cluster.cluster_id in endpoint.in_clusters
+            ):
+                endpoint.in_clusters.pop(self.cluster.cluster_id)
+                endpoint.add_input_cluster(self.cluster.cluster_id, self.cluster)
+            if (
+                ClusterType.Client in self.cluster_types
+                and self.cluster.cluster_id in endpoint.out_clusters
+            ):
+                endpoint.out_clusters.pop(self.cluster.cluster_id)
+                endpoint.add_output_cluster(self.cluster.cluster_id, self.cluster)
+
+
 @attrs.define(frozen=True, kw_only=True, repr=True)
 class EntityMetadata:
     """Metadata for an exposed entity."""
@@ -347,6 +376,9 @@ class QuirksV2RegistryEntry:
     adds_metadata: tuple[AddsMetadata] = attrs.field(factory=tuple)
     removes_metadata: tuple[RemovesMetadata] = attrs.field(factory=tuple)
     replaces_metadata: tuple[ReplacesMetadata] = attrs.field(factory=tuple)
+    replaces_cluster_occurrences_metadata: tuple[ReplaceClusterOccurrencesMetadata] = (
+        attrs.field(factory=tuple)
+    )
     entity_metadata: tuple[
         ZCLEnumMetadata
         | SwitchMetadata
@@ -388,6 +420,9 @@ def __init__(
         self.adds_metadata: list[AddsMetadata] = []
         self.removes_metadata: list[RemovesMetadata] = []
         self.replaces_metadata: list[ReplacesMetadata] = []
+        self.replaces_cluster_occurrences_metadata: list[
+            ReplaceClusterOccurrencesMetadata
+        ] = []
         self.entity_metadata: list[
             ZCLEnumMetadata
             | SwitchMetadata
@@ -539,6 +574,37 @@ def replaces(
         self.replaces_metadata.append(replace)
         return self
 
+    def replace_cluster_occurrences(
+        self,
+        replacement_cluster_class: type[Cluster | CustomCluster],
+        replace_server_instances: bool = True,
+        replace_client_instances: bool = True,
+    ) -> QuirkBuilder:
+        """Add a ReplaceClusterOccurrencesMetadata entry and returns self.
+
+        This method allows replacing a cluster on a device across all endpoints
+        for the specified cluster types when the quirk is applied.
+
+        replacement_cluster_class should be a subclass of Cluster or CustomCluster and
+        will be used to create a new cluster instance to replace the existing cluster.
+
+        If cluster_id is provided, it will be used as the cluster_id for the cluster to
+        be removed. If cluster_id is not provided, the cluster_id of the replacement
+        cluster will be used.
+        """
+        types = []
+        if replace_server_instances:
+            types.append(ClusterType.Server)
+        if replace_client_instances:
+            types.append(ClusterType.Client)
+        self.replaces_cluster_occurrences_metadata.append(
+            ReplaceClusterOccurrencesMetadata(  # type: ignore[call-arg]
+                cluster_types=types,
+                cluster=replacement_cluster_class,
+            )
+        )
+        return self
+
     def enum(
         self,
         attribute_name: str,
@@ -809,6 +875,9 @@ def add_to_registry(self) -> QuirksV2RegistryEntry:
             adds_metadata=tuple(self.adds_metadata),
             removes_metadata=tuple(self.removes_metadata),
             replaces_metadata=tuple(self.replaces_metadata),
+            replaces_cluster_occurrences_metadata=tuple(
+                self.replaces_cluster_occurrences_metadata
+            ),
             entity_metadata=tuple(self.entity_metadata),
             device_automation_triggers_metadata=self.device_automation_triggers_metadata,
         )