Repository: zigpy/zha-device-handlers
Pull request #3643 (2024-12-25 23:10:53)

Add support for Tuya DPs mapped to multiple cluster attributes
## Proposed change
<!--
  Explain your proposed change below.
-->
Some Tuya devices expose multiple atributes (like current+voltage+power) in a single DP. This PR adds support for easilly mapping those, including Builder support.


## Additional information
<!--
  Please include any additional information that is important to this PR.
  For example, if this PR is a potentially breaking change, mention that here.
  If this PR requires other PRs to be merged in HA Core or other projects, mention that.
  Lastly, if this PR fixes a specific issue, please include "Fixes #xxxx".
-->


## Checklist
<!--
  Put an 'x' in all boxes that apply.
  Note: You do not need to tick all boxes before creating a PR.
-->

- [x] The changes are tested and work correctly
- [x] `pre-commit` checks pass / the code has been formatted using Black
- [x] Tests have been added to verify that the new code works


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3643?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) Report
Attention: Patch coverage is `92.15686%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 89.82%. Comparing base [(`56001dd`)](https://app.codecov.io/gh/zigpy/zha-device-handlers/commit/56001dd08f154ed8446e09269162efd393b045b7?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) to head [(`f22ee00`)](https://app.codecov.io/gh/zigpy/zha-device-handlers/commit/f22ee00c6b5127aecabd02abc0113452745d1549?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).

| [Files with missing lines](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3643?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) | Patch % | Lines |
|---|---|---|
| [zhaquirks/tuya/builder/\_\_init\_\_.py](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3643?src=pr&el=tree&filepath=zhaquirks%2Ftuya%2Fbuilder%2F__init__.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy#diff-emhhcXVpcmtzL3R1eWEvYnVpbGRlci9fX2luaXRfXy5weQ==) | 50.00% | [3 Missing :warning: ](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3643?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) |
| [zhaquirks/tuya/\_\_init\_\_.py](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3643?src=pr&el=tree&filepath=zhaquirks%2Ftuya%2F__init__.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy#diff-emhhcXVpcmtzL3R1eWEvX19pbml0X18ucHk=) | 96.55% | [1 Missing :warning: ](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3643?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##              dev    #3643      +/-   ##
==========================================
- Coverage   89.84%   89.82%   -0.03%     
==========================================
  Files         322      322              
  Lines       10371    10378       +7     
==========================================
+ Hits         9318     9322       +4     
- Misses       1053     1056       +3     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3643?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).

jeverley: I was under the impression that this is already supported in v1 quirks, see the DPToAttributeMapping class:

```
@dataclasses.dataclass
class DPToAttributeMapping:
    """Container for datapoint to cluster attribute update mapping."""

    ep_attribute: str
    attribute_name: Union[str, tuple]
    converter: Optional[
        Callable[
            [
                Any,
            ],
            Any,
        ]
    ] = None
    endpoint_id: Optional[int] = None
```

`attribute_name` can be provided as a tuple containing multiple attribute names, this is then used in conjunction with a converter that returns a value for each attribute.

Below is an example that I used previously that calls a function in the TuyaPowerPhase class to extract voltage, current and power:

```
    dp_to_attribute: Dict[int, DPToAttributeMapping] = {
        TUYA_DP_POWER_PHASE: DPToAttributeMapping(
            TuyaElectricalMeasurement.ep_attribute,
            (
                "rms_voltage",
                "rms_current",
                "active_power",
            ),
            converter=lambda x: TuyaPowerPhase.variant_3(x),
        ),
    }

class TuyaPowerPhase:
    """Extract values from a Tuya power phase datapoint."""

    @staticmethod
    def variant_3(value) -> Tuple[t.uint_t, t.uint_t, int]:
        voltage = (value[0] << 8) | value[1]
        current = (value[2] << 16) | (value[3] << 8) | value[4]
        power = (value[5] << 16) | (value[6] << 8) | value[7]
        return voltage, current, power * 10
```

Before seeing this PR I as about to raise my own pull to amend the v2 tuya builder methods `tuya_dp` and `tuya_dp_attribute`  to support `str | tuple` for `attribute_name` to match the underlying DPToAttributeMapping class - are we overcomplicating things here?
abmantis: @jeverley TheJulianJES pointed that out too here: https://github.com/zigpy/zha-device-handlers/pull/3643#pullrequestreview-2574122550

I have missed that when I opened this. Still, this PR makes it a bit more explicit, and allows having a single DP mapped to different clusters, which is not possible with the old DPToAttributeMapping tuple approach.

Also, not all code is currently properly handling the tuple. See https://github.com/zigpy/zha-device-handlers/pull/3643#discussion_r1929645872

The plan is to check if we can drop the old tuple after this is merge (and migrate any quirks using it).
TheJulianJES: And we should still remove the old "multiple DP" way in another PR.

Patch:
diff --git a/zhaquirks/tuya/__init__.py b/zhaquirks/tuya/__init__.py
index 2fe2ceb..0c6ffcf 100644
--- a/zhaquirks/tuya/__init__.py
+++ b/zhaquirks/tuya/__init__.py
@@ -1533,35 +1533,43 @@ class TuyaNewManufCluster(CustomCluster):
         ),
     }
 
-    dp_to_attribute: dict[int, DPToAttributeMapping] = {}
+    dp_to_attribute: dict[int, DPToAttributeMapping | list[DPToAttributeMapping]] = {}
     data_point_handlers: dict[int, str] = {}
 
     def __init__(self, *args, **kwargs):
         """Initialize the cluster and mark attributes as valid on LocalDataClusters."""
         super().__init__(*args, **kwargs)
-        for dp_map in self.dp_to_attribute.values():
+
+        self._dp_to_attributes: dict[int, list[DPToAttributeMapping]] = {
+            dp: attr if isinstance(attr, list) else [attr]
+            for dp, attr in self.dp_to_attribute.items()
+        }
+        for dp_map in self._dp_to_attributes.values():
             # get the endpoint that is being mapped to
             endpoint = self.endpoint
-            if dp_map.endpoint_id:
-                endpoint = self.endpoint.device.endpoints.get(dp_map.endpoint_id)
-
-            # the endpoint to be mapped to might not actually exist within all quirks
-            if not endpoint:
-                continue
-
-            cluster = getattr(endpoint, dp_map.ep_attribute, None)
-            # the cluster to be mapped to might not actually exist within all quirks
-            if not cluster:
-                continue
-
-            # mark mapped to attribute as valid if existing and if on a LocalDataCluster
-            attr = cluster.attributes_by_name.get(dp_map.attribute_name)
-            if attr and isinstance(cluster, LocalDataCluster):
-                # _VALID_ATTRIBUTES is only a class variable, but as want to modify it
-                # per instance here, we need to create an instance variable first
-                if "_VALID_ATTRIBUTES" not in cluster.__dict__:
-                    cluster._VALID_ATTRIBUTES = set()
-                cluster._VALID_ATTRIBUTES.add(attr.id)
+            for mapped_attr in dp_map:
+                if mapped_attr.endpoint_id:
+                    endpoint = self.endpoint.device.endpoints.get(
+                        mapped_attr.endpoint_id
+                    )
+
+                # the endpoint to be mapped to might not actually exist within all quirks
+                if not endpoint:
+                    continue
+
+                cluster = getattr(endpoint, mapped_attr.ep_attribute, None)
+                # the cluster to be mapped to might not actually exist within all quirks
+                if not cluster:
+                    continue
+
+                # mark mapped to attribute as valid if existing and if on a LocalDataCluster
+                attr = cluster.attributes_by_name.get(mapped_attr.attribute_name)
+                if attr and isinstance(cluster, LocalDataCluster):
+                    # _VALID_ATTRIBUTES is only a class variable, but as want to modify it
+                    # per instance here, we need to create an instance variable first
+                    if "_VALID_ATTRIBUTES" not in cluster.__dict__:
+                        cluster._VALID_ATTRIBUTES = set()
+                    cluster._VALID_ATTRIBUTES.add(attr.id)
 
     def handle_cluster_request(
         self,
@@ -1640,27 +1648,29 @@ class TuyaNewManufCluster(CustomCluster):
     def _dp_2_attr_update(self, datapoint: TuyaDatapointData) -> None:
         """Handle data point to attribute report conversion."""
         try:
-            dp_map = self.dp_to_attribute[datapoint.dp]
+            dp_map = self._dp_to_attributes[datapoint.dp]
         except KeyError:
             self.debug("No attribute mapping for %s data point", datapoint.dp)
             return
 
         endpoint = self.endpoint
-        if dp_map.endpoint_id:
-            endpoint = self.endpoint.device.endpoints[dp_map.endpoint_id]
-        cluster = getattr(endpoint, dp_map.ep_attribute)
-        value = datapoint.data.payload
-        if dp_map.converter:
-            value = dp_map.converter(value)
-
-        if isinstance(dp_map.attribute_name, tuple):
-            for k, v in zip(dp_map.attribute_name, value):
-                if isinstance(v, AttributeWithMask):
-                    v = cluster.get(k, 0) & (~v.mask) | v.value
-                cluster.update_attribute(k, v)
-        else:
-            if isinstance(value, AttributeWithMask):
-                value = (
-                    cluster.get(dp_map.attribute_name, 0) & (~value.mask) | value.value
-                )
-            cluster.update_attribute(dp_map.attribute_name, value)
+        for mapped_attr in dp_map:
+            if mapped_attr.endpoint_id:
+                endpoint = self.endpoint.device.endpoints[mapped_attr.endpoint_id]
+            cluster = getattr(endpoint, mapped_attr.ep_attribute)
+            value = datapoint.data.payload
+            if mapped_attr.converter:
+                value = mapped_attr.converter(value)
+
+            if isinstance(mapped_attr.attribute_name, tuple):
+                for k, v in zip(mapped_attr.attribute_name, value):
+                    if isinstance(v, AttributeWithMask):
+                        v = cluster.get(k, 0) & (~v.mask) | v.value
+                    cluster.update_attribute(k, v)
+            else:
+                if isinstance(value, AttributeWithMask):
+                    value = (
+                        cluster.get(mapped_attr.attribute_name, 0) & (~value.mask)
+                        | value.value
+                    )
+                cluster.update_attribute(mapped_attr.attribute_name, value)
diff --git a/zhaquirks/tuya/builder/__init__.py b/zhaquirks/tuya/builder/__init__.py
index d9e921d..475328b 100644
--- a/zhaquirks/tuya/builder/__init__.py
+++ b/zhaquirks/tuya/builder/__init__.py
@@ -196,7 +196,7 @@ class TuyaQuirkBuilder(QuirkBuilder):
     ) -> None:
         """Init the TuyaQuirkBuilder."""
         self.tuya_data_point_handlers: dict[int, str] = {}
-        self.tuya_dp_to_attribute: dict[int, DPToAttributeMapping] = {}
+        self.tuya_dp_to_attribute: dict[int, list[DPToAttributeMapping]] = {}
         self.new_attributes: set[foundation.ZCLAttributeDef] = set()
         super().__init__(manufacturer, model, registry)
         # quirk_file will point to the init call above if called from this QuirkBuilder,
@@ -503,20 +503,33 @@ class TuyaQuirkBuilder(QuirkBuilder):
     ) -> QuirkBuilder:  # fmt: skip
         """Add Tuya DP Converter."""
 
-        if dp_id in self.tuya_dp_to_attribute:
-            raise ValueError(f"DP {dp_id} is already mapped.")
-
-        self.tuya_dp_to_attribute.update(
-            {
-                dp_id: DPToAttributeMapping(
+        self.tuya_dp_multi(
+            dp_id,
+            [
+                DPToAttributeMapping(
                     ep_attribute,
                     attribute_name,
                     converter=converter,
                     dp_converter=dp_converter,
                     endpoint_id=endpoint_id,
                 )
-            }
+            ],
+            dp_handler,
         )
+        return self
+
+    def tuya_dp_multi(
+        self,
+        dp_id: int,
+        attribute_mapping: list[DPToAttributeMapping],
+        dp_handler: str = "_dp_2_attr_update",
+    ) -> QuirkBuilder:  # fmt: skip
+        """Add Tuya DP Converter that maps to multiple attributes."""
+
+        if dp_id in self.tuya_dp_to_attribute:
+            raise ValueError(f"DP {dp_id} is already mapped.")
+
+        self.tuya_dp_to_attribute.update({dp_id: attribute_mapping})
         self.tuya_data_point_handlers.update({dp_id: dp_handler})
         return self
 
@@ -814,7 +827,7 @@ class TuyaQuirkBuilder(QuirkBuilder):
                 """Replacement Tuya Cluster."""
 
                 data_point_handlers: dict[int, str]
-                dp_to_attribute: dict[int, DPToAttributeMapping]
+                dp_to_attribute: dict[int, list[DPToAttributeMapping]]
 
                 class AttributeDefs(NewAttributeDefs):
                     """Attribute Definitions."""
diff --git a/zhaquirks/tuya/mcu/__init__.py b/zhaquirks/tuya/mcu/__init__.py
index 26a11fa..2dfd104 100644
--- a/zhaquirks/tuya/mcu/__init__.py
+++ b/zhaquirks/tuya/mcu/__init__.py
@@ -288,25 +288,26 @@ class TuyaMCUCluster(TuyaAttributesCluster, TuyaNewManufCluster):
     def get_dp_mapping(
         self, endpoint_id: int, attribute_name: str
     ) -> Optional[tuple[int, DPToAttributeMapping]]:
-        """Search for the DP in dp_to_attribute."""
+        """Search for the DP in _dp_to_attributes."""
 
         result = {}
-        for dp, dp_mapping in self.dp_to_attribute.items():
-            if (
-                attribute_name == dp_mapping.attribute_name
-                or (
-                    isinstance(dp_mapping.attribute_name, tuple)
-                    and attribute_name in dp_mapping.attribute_name
-                )
-            ) and (
-                (
-                    dp_mapping.endpoint_id is None
-                    and endpoint_id == self.endpoint.endpoint_id
-                )
-                or (endpoint_id == dp_mapping.endpoint_id)
-            ):
-                self.debug("get_dp_mapping --> found DP: %s", dp)
-                result[dp] = dp_mapping
+        for dp, dp_mapping in self._dp_to_attributes.items():
+            for mapped_attr in dp_mapping:
+                if (
+                    attribute_name == mapped_attr.attribute_name
+                    or (
+                        isinstance(mapped_attr.attribute_name, tuple)
+                        and attribute_name in mapped_attr.attribute_name
+                    )
+                ) and (
+                    (
+                        mapped_attr.endpoint_id is None
+                        and endpoint_id == self.endpoint.endpoint_id
+                    )
+                    or (endpoint_id == mapped_attr.endpoint_id)
+                ):
+                    self.debug("get_dp_mapping --> found DP: %s", dp)
+                    result[dp] = mapped_attr
         return result
 
     def handle_mcu_version_response(self, payload: MCUVersion) -> foundation.Status:


Test patch:
diff --git a/tests/test_tuya_builder.py b/tests/test_tuya_builder.py
index b223b97..8041ae2 100644
--- a/tests/test_tuya_builder.py
+++ b/tests/test_tuya_builder.py
@@ -1,5 +1,6 @@
 """Tests for TuyaQuirkBuilder."""
 
+from collections.abc import ByteString
 import datetime
 from unittest import mock
 
@@ -9,6 +10,7 @@ from zigpy.quirks.v2 import CustomDeviceV2
 import zigpy.types as t
 from zigpy.zcl import foundation
 from zigpy.zcl.clusters.general import Basic
+from zigpy.zcl.clusters.homeautomation import ElectricalMeasurement
 
 from tests.common import ClusterListener, MockDatetime, wait_for_zigpy_tasks
 import zhaquirks
@@ -16,6 +18,11 @@ from zhaquirks.const import BatterySize
 from zhaquirks.tuya import (
     TUYA_QUERY_DATA,
     TUYA_SET_TIME,
+    DPToAttributeMapping,
+    TuyaCommand,
+    TuyaData,
+    TuyaDatapointData,
+    TuyaLocalCluster,
     TuyaPowerConfigurationCluster,
     TuyaPowerConfigurationCluster2AAA,
 )
@@ -156,6 +163,18 @@ async def test_tuya_quirkbuilder(device_mock):
     class ModTuyaMCUCluster(TuyaMCUCluster):
         """Modified Cluster."""
 
+    class Tuya3PhaseElectricalMeasurement(ElectricalMeasurement, TuyaLocalCluster):
+        """Tuya Electrical Measurement cluster."""
+
+    def dpToPower(data: ByteString) -> int:
+        return data[0]
+
+    def dpToCurrent(data: ByteString) -> int:
+        return data[1]
+
+    def dpToVoltage(data: ByteString) -> int:
+        return data[2]
+
     entry = (
         TuyaQuirkBuilder(device_mock.manufacturer, device_mock.model, registry=registry)
         .tuya_battery(dp_id=1)
@@ -193,6 +212,27 @@ async def test_tuya_quirkbuilder(device_mock):
             translation_key="test_enum",
             fallback_name="Test enum",
         )
+        .tuya_dp_multi(
+            dp_id=11,
+            attribute_mapping=[
+                DPToAttributeMapping(
+                    ep_attribute=Tuya3PhaseElectricalMeasurement.ep_attribute,
+                    attribute_name="active_power",
+                    converter=dpToPower,
+                ),
+                DPToAttributeMapping(
+                    ep_attribute=Tuya3PhaseElectricalMeasurement.ep_attribute,
+                    attribute_name="rms_current",
+                    converter=dpToCurrent,
+                ),
+                DPToAttributeMapping(
+                    ep_attribute=Tuya3PhaseElectricalMeasurement.ep_attribute,
+                    attribute_name="rms_voltage",
+                    converter=dpToVoltage,
+                ),
+            ],
+        )
+        .adds(Tuya3PhaseElectricalMeasurement)
         .skip_configuration()
         .add_to_registry(replacement_cluster=ModTuyaMCUCluster)
     )
@@ -251,6 +291,17 @@ async def test_tuya_quirkbuilder(device_mock):
     assert tuya_listener.attribute_updates[0][0] == 0xEF0A
     assert tuya_listener.attribute_updates[0][1] == TestEnum.B
 
+    electric_data = TuyaCommand(
+        status=0,
+        tsn=2,
+        datapoints=[TuyaDatapointData(11, TuyaData("345"))],
+    )
+    tuya_cluster.handle_get_data(electric_data)
+    electrical_meas_cluster = ep.electrical_measurement
+    assert electrical_meas_cluster.get("active_power") == "3"
+    assert electrical_meas_cluster.get("rms_current") == "4"
+    assert electrical_meas_cluster.get("rms_voltage") == "5"
+
 
 async def test_tuya_quirkbuilder_duplicated_mappings(device_mock):
     """Test that mapping the same DP multiple times will raise."""
@@ -268,6 +319,24 @@ async def test_tuya_quirkbuilder_duplicated_mappings(device_mock):
             .add_to_registry()
         )
 
+    with pytest.raises(ValueError):
+        (
+            TuyaQuirkBuilder(
+                device_mock.manufacturer, device_mock.model, registry=registry
+            )
+            .tuya_battery(dp_id=1)
+            .tuya_dp_multi(
+                dp_id=1,
+                attribute_mapping=[
+                    DPToAttributeMapping(
+                        ep_attribute=ElectricalMeasurement.ep_attribute,
+                        attribute_name="active_power",
+                    ),
+                ],
+            )
+            .add_to_registry()
+        )
+
 
 @pytest.mark.parametrize(
     "read_attr_spell,data_query_spell",