Improve typing in Tuya builder (#4302)

Author: epenet

zhaquirks/tuya/builder/__init__.py

@@ -6,7 +6,7 @@
 import math
 import pathlib
 from types import FrameType
-from typing import Any
+from typing import Any, Self
 
 from zigpy.quirks import _DEVICE_REGISTRY
 from zigpy.quirks.registry import DeviceRegistry
@@ -217,7 +217,7 @@ def _tuya_battery(
         dp_id: int,
         power_cfg: PowerConfiguration,
         scale: float,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Battery Power Configuration."""
         self.tuya_dp(
             dp_id,
@@ -236,7 +236,7 @@ def tuya_battery(
         battery_qty: int | None = 2,
         battery_voltage: int | None = None,
         scale: float = 2,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Battery Power Configuration."""
 
         if power_cfg:
@@ -265,7 +265,7 @@ def tuya_illuminance(
         converter: Callable[[Any], Any] | None = (
             lambda x: 10000 * math.log10(x) + 1 if x != 0 else 0
         ),
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Illuminance Configuration."""
         self.tuya_dp(
             dp_id,
@@ -276,7 +276,7 @@ def tuya_illuminance(
         self.adds(illuminance_cfg)
         return self
 
-    def tuya_contact(self, dp_id: int):
+    def tuya_contact(self, dp_id: int) -> Self:
         """Add a Tuya IAS contact sensor."""
         self.tuya_ias(
             dp_id=dp_id,
@@ -290,7 +290,7 @@ def tuya_co2(
         dp_id: int,
         co2_cfg: TuyaLocalCluster = TuyaCO2Concentration,
         scale: float = 1e-6,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya CO2 Configuration."""
         self.tuya_dp(
             dp_id,
@@ -306,7 +306,7 @@ def tuya_electrical_conductivity(
         dp_id: int,
         ec_cfg: TuyaLocalCluster = TuyaElectricalConductivity,
         scale: float = 1,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Electrical Conductivity Configuration."""
         self.tuya_dp(
             dp_id,
@@ -326,7 +326,7 @@ def tuya_formaldehyde(
             ((MOL_VOL_AIR_NTP * x) / TuyaFormaldehydeConcentration.MOLECULAR_MASS), 2
         )
         * 1e-6,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Formaldehyde Configuration."""
         self.tuya_dp(
             dp_id,
@@ -342,7 +342,7 @@ def tuya_pm25(
         dp_id: int,
         pm25_cfg: TuyaLocalCluster = TuyaPM25Concentration,
         scale: float = 1,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya PM25 Configuration."""
         self.tuya_dp(
             dp_id,
@@ -353,7 +353,7 @@ def tuya_pm25(
         self.adds(pm25_cfg)
         return self
 
-    def tuya_gas(self, dp_id: int):
+    def tuya_gas(self, dp_id: int) -> Self:
         """Add a Tuya IAS gas sensor."""
         self.tuya_ias(
             dp_id=dp_id,
@@ -362,7 +362,7 @@ def tuya_gas(self, dp_id: int):
         )
         return self
 
-    def tuya_smoke(self, dp_id: int):
+    def tuya_smoke(self, dp_id: int) -> Self:
         """Add a Tuya IAS smoke/fire sensor."""
         self.tuya_ias(
             dp_id=dp_id,
@@ -376,7 +376,7 @@ def tuya_ias(
         dp_id: int,
         ias_cfg: TuyaLocalCluster,
         converter: Callable[[Any], Any] | None = None,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya IAS Configuration."""
         self.tuya_dp(
             dp_id,
@@ -392,7 +392,7 @@ def tuya_metering(
         dp_id: int,
         metering_cfg: TuyaLocalCluster = TuyaValveWaterConsumedNoInstDemand,
         scale: float = 1,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Metering Configuration."""
         self.tuya_dp(
             dp_id,
@@ -407,7 +407,7 @@ def tuya_onoff(
         self,
         dp_id: int,
         onoff_cfg: TuyaLocalCluster = TuyaOnOffNM,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya OnOff Configuration."""
         self.tuya_dp(
             dp_id,
@@ -422,7 +422,7 @@ def tuya_humidity(
         dp_id: int,
         rh_cfg: TuyaLocalCluster = TuyaRelativeHumidity,
         scale: float = 100,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Relative Humidity Configuration."""
         self.tuya_dp(
             dp_id,
@@ -438,7 +438,7 @@ def tuya_soil_moisture(
         dp_id: int,
         soil_cfg: TuyaLocalCluster = TuyaSoilMoisture,
         scale: float = 100,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Soil Moisture Configuration."""
         self.tuya_dp(
             dp_id,
@@ -454,7 +454,7 @@ def tuya_temperature(
         dp_id: int,
         temp_cfg: TuyaLocalCluster = TuyaTemperatureMeasurement,
         scale: float = 100,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya Temperature Configuration."""
         self.tuya_dp(
             dp_id,
@@ -465,7 +465,7 @@ def tuya_temperature(
         self.adds(temp_cfg)
         return self
 
-    def tuya_vibration(self, dp_id: int):
+    def tuya_vibration(self, dp_id: int) -> Self:
         """Add a Tuya IAS vibration sensor."""
         self.tuya_ias(
             dp_id=dp_id,
@@ -479,7 +479,7 @@ def tuya_voc(
         dp_id: int,
         voc_cfg: TuyaLocalCluster = TuyaAirQualityVOC,
         scale: float = 1e-6,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add a Tuya VOC Configuration."""
         self.tuya_dp(
             dp_id,
@@ -497,7 +497,7 @@ def tuya_attribute(
         type: type = t.uint16_t,
         access: foundation.ZCLAttributeAccess = foundation.ZCLAttributeAccess.NONE,
         is_manufacturer_specific=True,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add an attribute to AttributeDefs."""
         attr_id: int = int.from_bytes([0xEF, dp_id])
 
@@ -522,7 +522,7 @@ def tuya_dp(
         dp_converter: Callable[[Any], Any] | None = None,
         endpoint_id: int | None = None,
         dp_handler: str = "_dp_2_attr_update",
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add Tuya DP Converter."""
 
         self.tuya_dp_multi(
@@ -545,7 +545,7 @@ def tuya_dp_multi(
         dp_id: int,
         attribute_mapping: list[DPToAttributeMapping],
         dp_handler: str = "_dp_2_attr_update",
-    ) -> QuirkBuilder:  # fmt: skip
+    ) -> Self:  # fmt: skip
         """Add Tuya DP Converter that maps to multiple attributes."""
 
         if dp_id in self.tuya_dp_to_attribute:
@@ -567,7 +567,7 @@ def tuya_dp_attribute(
         type: type = t.uint16_t,
         access: foundation.ZCLAttributeAccess = foundation.ZCLAttributeAccess.NONE,
         is_manufacturer_specific=True,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add an Tuya DataPoint and corresponding AttributeDef."""
         self.tuya_attribute(
             dp_id=dp_id,
@@ -602,7 +602,7 @@ def tuya_switch(
         attribute_initialized_from_cache: bool = True,
         translation_key: str | None = None,
         fallback_name: str | None = None,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add an EntityMetadata containing SwitchMetadata and return self.
 
         This method allows exposing a switch entity in Home Assistant.
@@ -645,7 +645,7 @@ def tuya_enum(
         attribute_initialized_from_cache: bool = True,
         translation_key: str | None = None,
         fallback_name: str | None = None,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add an EntityMetadata containing ZCLEnumMetadata and return self.
 
         This method allows exposing an enum based entity in Home Assistant.
@@ -691,7 +691,7 @@ def tuya_number(
         attribute_initialized_from_cache: bool = True,
         translation_key: str | None = None,
         fallback_name: str | None = None,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add an EntityMetadata containing NumberMetadata and return self.
 
         This method allows exposing a number entity in Home Assistant.
@@ -733,7 +733,7 @@ def tuya_binary_sensor(
         attribute_initialized_from_cache: bool = True,
         translation_key: str | None = None,
         fallback_name: str | None = None,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add an EntityMetadata containing BinarySensorMetadata and return self.
 
         This method allows exposing a binary sensor entity in Home Assistant.
@@ -777,7 +777,7 @@ def tuya_sensor(
         attribute_initialized_from_cache: bool = True,
         translation_key: str | None = None,
         fallback_name: str | None = None,
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Add an EntityMetadata containing ZCLSensorMetadata and return self.
 
         This method allows exposing a sensor entity in Home Assistant.
@@ -812,7 +812,7 @@ def tuya_sensor(
 
     def tuya_enchantment(
         self, read_attr_spell: bool = True, data_query_spell: bool = False
-    ) -> QuirkBuilder:
+    ) -> Self:
         """Set the Tuya enchantment spells."""
 
         class EnchantedDeviceV2(CustomDeviceV2, BaseEnchantedDevice):

zhaquirks/tuya/ts0601_power.py

@@ -12,8 +12,9 @@
 from zigpy.zcl.clusters.general import LevelControl, OnOff
 from zigpy.zcl.clusters.homeautomation import ElectricalMeasurement
 
-from zhaquirks.tuya import DPToAttributeMapping, TuyaLocalCluster
+from zhaquirks.tuya import TuyaLocalCluster
 from zhaquirks.tuya.builder import TuyaQuirkBuilder
+from zhaquirks.tuya.mcu import DPToAttributeMapping
 
 
 def dp_to_power(data: bytes) -> int: