Allow quirks v2 to specify primary entities (#468)

* Allow quirks v2 to specify primary entities

* Add a unit test

Author: puddly

tests/test_device.py

@@ -13,6 +13,7 @@
 from zigpy.quirks.v2 import DeviceAlertLevel, DeviceAlertMetadata, QuirkBuilder
 import zigpy.types
 from zigpy.zcl.clusters import general
+from zigpy.zcl.clusters.general import PowerConfiguration
 from zigpy.zcl.foundation import Status, WriteAttributesResponse
 import zigpy.zdo.types as zdo_t
 
@@ -895,6 +896,35 @@ async def test_primary_entity_computation(
         ]
 
 
+async def test_quirks_v2_primary_entity(zha_gateway: Gateway) -> None:
+    """Test quirks v2 primary entity."""
+    registry = DeviceRegistry()
+
+    (
+        QuirkBuilder("CentraLite", "3405-L", registry=registry)
+        .sensor(
+            attribute_name=PowerConfiguration.AttributeDefs.battery_quantity.id,
+            cluster_id=PowerConfiguration.cluster_id,
+            translation_key="battery_quantity",
+            fallback_name="Battery quantity",
+            primary=True,
+        )
+        .add_to_registry()
+    )
+
+    zigpy_dev = registry.get_device(
+        await zigpy_device_from_json(
+            zha_gateway.application_controller,
+            "tests/data/devices/centralite-3405-l.json",
+        )
+    )
+
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    (primary,) = [e for e in zha_device.platform_entities.values() if e.primary]
+    assert primary.translation_key == "battery_quantity"
+
+
 async def test_quirks_v2_prevent_default_entities(zha_gateway: Gateway) -> None:
     """Test quirks v2 can prevent creating default entities."""
     registry = DeviceRegistry()

zha/application/platforms/__init__.py

@@ -123,7 +123,7 @@ class BaseEntity(LogMixin, EventBase):
     _attr_state_class: str | None = None
     _attr_enabled: bool = True
     _attr_always_supported: bool = False
-    _attr_primary: bool = False
+    _attr_primary: bool | None = None
 
     # When two entities both want to be primary, the one with the higher weight will be
     # chosen. If there is a tie, both lose.
@@ -173,10 +173,13 @@ def enabled(self, value: bool) -> None:
     @property
     def primary(self) -> bool:
         """Return if the entity is the primary device control."""
+        if self._attr_primary is None:
+            return False
+
         return self._attr_primary
 
     @primary.setter
-    def primary(self, value: bool) -> None:
+    def primary(self, value: bool | None) -> None:
         """Set the entity as the primary device control."""
         self._attr_primary = value
 
@@ -404,6 +407,9 @@ def _init_from_quirks_metadata(self, entity_metadata: EntityMetadata) -> None:
         else:
             self._attr_entity_category = None
 
+        if entity_metadata.primary is not None:
+            self._attr_primary = entity_metadata.primary
+
     @cached_property
     def identifiers(self) -> PlatformEntityIdentifiers:
         """Return a dict with the information necessary to identify this entity."""

zha/zigbee/device.py

@@ -1207,11 +1207,27 @@ def log(self, level: int, msg: str, *args: Any, **kwargs: Any) -> None:
     def _compute_primary_entity(self) -> None:
         """Compute the primary entity for this device."""
 
-        # Only consider non-counter entities
+        # First, check if any entity is explicitly primary
+        explicitly_primary = [
+            entity for entity in self._platform_entities.values() if entity.primary
+        ]
+
+        if len(explicitly_primary) == 1:
+            self.debug(
+                "Device has a single explicitly primary entity,"
+                " not performing weight matching"
+            )
+            return
+
+        # It should not be possible for there to be more than one
+        assert not explicitly_primary
+
+        # For weight matching, only consider non-counter entities and entities which are
+        # not explicitly marked as not primary
         candidates = [
             e
             for e in self._platform_entities.values()
-            if e.enabled and hasattr(e, "info_object")
+            if e.enabled and hasattr(e, "info_object") and e._attr_primary is not False
         ]
         candidates.sort(reverse=True, key=lambda e: e.primary_weight)