Implement quirks v2 attribute_converter (#360)

Author: TheJulianJES

tests/test_binary_sensor.py

@@ -4,8 +4,12 @@
 from unittest.mock import MagicMock, call
 
 import pytest
+from zigpy.profiles import zha
 import zigpy.profiles.zha
+from zigpy.quirks import DeviceRegistry
+from zigpy.quirks.v2 import CustomDeviceV2, QuirkBuilder
 from zigpy.zcl.clusters import general, measurement, security
+from zigpy.zcl.clusters.general import OnOff
 
 from tests.common import (
     SIG_EP_INPUT,
@@ -22,7 +26,12 @@
 from zha.application import Platform
 from zha.application.gateway import Gateway
 from zha.application.platforms import PlatformEntity
-from zha.application.platforms.binary_sensor import Accelerometer, IASZone, Occupancy
+from zha.application.platforms.binary_sensor import (
+    Accelerometer,
+    BinarySensor,
+    IASZone,
+    Occupancy,
+)
 from zha.zigbee.cluster_handlers.const import SMARTTHINGS_ACCELERATION_CLUSTER
 
 DEVICE_IAS = {
@@ -201,3 +210,49 @@ async def test_smarttthings_multi(
             {"attribute_id": 18, "attribute_name": "x_axis", "attribute_value": 120},
         )
     ]
+
+
+async def test_quirks_binary_sensor_attr_converter(zha_gateway: Gateway) -> None:
+    """Test ZHA quirks v2 binary_sensor with attribute_converter."""
+
+    registry = DeviceRegistry()
+    zigpy_dev = create_mock_zigpy_device(
+        zha_gateway,
+        {
+            1: {
+                SIG_EP_INPUT: [general.Basic.cluster_id, general.OnOff.cluster_id],
+                SIG_EP_OUTPUT: [],
+                SIG_EP_TYPE: zha.DeviceType.SIMPLE_SENSOR,
+            }
+        },
+        manufacturer="manufacturer",
+        model="model",
+    )
+
+    (
+        QuirkBuilder(zigpy_dev.manufacturer, zigpy_dev.model, registry=registry)
+        .binary_sensor(
+            OnOff.AttributeDefs.on_off.name,
+            OnOff.cluster_id,
+            translation_key="on_off",
+            fallback_name="On/off",
+            attribute_converter=lambda x: not bool(x),  # invert value with lambda
+        )
+        .add_to_registry()
+    )
+
+    zigpy_device_ = registry.get_device(zigpy_dev)
+
+    assert isinstance(zigpy_device_, CustomDeviceV2)
+    cluster = zigpy_device_.endpoints[1].on_off
+
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_device_)
+    entity = get_entity(zha_device, platform=Platform.BINARY_SENSOR)
+    assert isinstance(entity, BinarySensor)
+
+    # send updated value, check if the value is inverted
+    await send_attributes_report(zha_gateway, cluster, {"on_off": 1})
+    assert entity.is_on is False
+
+    await send_attributes_report(zha_gateway, cluster, {"on_off": 0})
+    assert entity.is_on is True

tests/test_sensor.py

@@ -11,15 +11,17 @@
 import pytest
 from zhaquirks.danfoss import thermostat as danfoss_thermostat
 from zigpy.device import Device as ZigpyDevice
+from zigpy.profiles import zha
 import zigpy.profiles.zha
-from zigpy.quirks import CustomCluster, get_device
+from zigpy.quirks import CustomCluster, DeviceRegistry, get_device
 from zigpy.quirks.v2 import CustomDeviceV2, QuirkBuilder, ReportingConfig
 from zigpy.quirks.v2.homeassistant.sensor import (
     SensorDeviceClass as SensorDeviceClassV2,
 )
 import zigpy.types as t
 from zigpy.zcl import Cluster
 from zigpy.zcl.clusters import general, homeautomation, hvac, measurement, smartenergy
+from zigpy.zcl.clusters.general import AnalogInput
 from zigpy.zcl.clusters.manufacturer_specific import ManufacturerSpecificCluster
 
 from tests.common import (
@@ -1679,3 +1681,51 @@ async def test_danfoss_thermostat_sw_error(zha_gateway: Gateway) -> None:
     assert entity.extra_state_attribute_names
     assert "Top_pcb_sensor_error" in entity.extra_state_attribute_names
     assert entity.state["Top_pcb_sensor_error"]
+
+
+async def test_quirks_sensor_attr_converter(zha_gateway: Gateway) -> None:
+    """Test ZHA quirks v2 sensor with attribute_converter."""
+
+    registry = DeviceRegistry()
+    zigpy_dev = create_mock_zigpy_device(
+        zha_gateway,
+        {
+            1: {
+                SIG_EP_INPUT: [
+                    general.Basic.cluster_id,
+                    general.AnalogInput.cluster_id,
+                ],
+                SIG_EP_OUTPUT: [],
+                SIG_EP_TYPE: zha.DeviceType.SIMPLE_SENSOR,
+            }
+        },
+        manufacturer="manufacturer",
+        model="model",
+    )
+
+    (
+        QuirkBuilder(zigpy_dev.manufacturer, zigpy_dev.model, registry=registry)
+        .sensor(
+            AnalogInput.AttributeDefs.present_value.name,
+            AnalogInput.cluster_id,
+            translation_key="quirks_sensor",
+            fallback_name="Quirks sensor",
+            attribute_converter=lambda x: x + 100,
+        )
+        .add_to_registry()
+    )
+
+    zigpy_device_ = registry.get_device(zigpy_dev)
+
+    assert isinstance(zigpy_device_, CustomDeviceV2)
+    cluster = zigpy_device_.endpoints[1].analog_input
+
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_device_)
+    entity = get_entity(zha_device, platform=Platform.SENSOR, qualifier="present_value")
+
+    # send updated value, check if the value is converted
+    await send_attributes_report(zha_gateway, cluster, {"present_value": 100})
+    assert entity.state["state"] == 200.0
+
+    await send_attributes_report(zha_gateway, cluster, {"present_value": 0})
+    assert entity.state["state"] == 100.0

zha/application/platforms/binary_sensor/__init__.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 import functools
 import logging
+import typing
 from typing import TYPE_CHECKING
 
 from zhaquirks.quirk_ids import DANFOSS_ALLY_THERMOSTAT
@@ -59,6 +60,7 @@ class BinarySensor(PlatformEntity):
 
     _attr_device_class: BinarySensorDeviceClass | None
     _attribute_name: str
+    _attribute_converter: typing.Callable[[typing.Any], typing.Any] | None = None
     PLATFORM: Platform = Platform.BINARY_SENSOR
 
     def __init__(
@@ -82,6 +84,8 @@ def _init_from_quirks_metadata(self, entity_metadata: BinarySensorMetadata) -> N
         """Init this entity from the quirks metadata."""
         super()._init_from_quirks_metadata(entity_metadata)
         self._attribute_name = entity_metadata.attribute_name
+        if entity_metadata.attribute_converter is not None:
+            self._attribute_converter = entity_metadata.attribute_converter
         if entity_metadata.device_class is not None:
             self._attr_device_class = validate_device_class(
                 BinarySensorDeviceClass,
@@ -113,6 +117,8 @@ def is_on(self) -> bool:
         )
         if raw_state is None:
             return False
+        if self._attribute_converter:
+            return self._attribute_converter(raw_state)
         return self.parse(raw_state)
 
     def handle_cluster_handler_attribute_updated(

zha/application/platforms/sensor/__init__.py

@@ -9,6 +9,7 @@
 import functools
 import logging
 import numbers
+import typing
 from typing import TYPE_CHECKING, Any, Self
 
 from zhaquirks.danfoss import thermostat as danfoss_thermostat
@@ -150,6 +151,7 @@ class Sensor(PlatformEntity):
 
     PLATFORM = Platform.SENSOR
     _attribute_name: int | str | None = None
+    _attribute_converter: typing.Callable[[typing.Any], typing.Any] | None = None
     _decimals: int = 1
     _divisor: int = 1
     _multiplier: int | float = 1
@@ -226,6 +228,8 @@ def _init_from_quirks_metadata(self, entity_metadata: ZCLSensorMetadata) -> None
         """Init this entity from the quirks metadata."""
         super()._init_from_quirks_metadata(entity_metadata)
         self._attribute_name = entity_metadata.attribute_name
+        if entity_metadata.attribute_converter is not None:
+            self._attribute_converter = entity_metadata.attribute_converter
         if entity_metadata.divisor is not None:
             self._divisor = entity_metadata.divisor
         if entity_metadata.multiplier is not None:
@@ -275,6 +279,8 @@ def native_value(self) -> date | datetime | str | int | float | None:
         raw_state = self._cluster_handler.cluster.get(self._attribute_name)
         if raw_state is None:
             return None
+        if self._attribute_converter:
+            return self._attribute_converter(raw_state)
         return self.formatter(raw_state)
 
     def handle_cluster_handler_attribute_updated(