Migrate Tuya climate (swing) to use wrapper class (home-assistant#157646)

Author: epenet

homeassistant/components/tuya/climate.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Self
 
 from tuya_sharing import CustomerDevice, Manager
 
@@ -32,7 +32,12 @@
     DPCode,
 )
 from .entity import TuyaEntity
-from .models import DPCodeBooleanWrapper, DPCodeEnumWrapper, DPCodeIntegerWrapper
+from .models import (
+    DeviceWrapper,
+    DPCodeBooleanWrapper,
+    DPCodeEnumWrapper,
+    DPCodeIntegerWrapper,
+)
 
 TUYA_HVAC_TO_HA = {
     "auto": HVACMode.HEAT_COOL,
@@ -56,6 +61,84 @@ def read_device_status(self, device: CustomerDevice) -> int | None:
         return round(value)
 
 
+@dataclass(kw_only=True)
+class _SwingModeWrapper(DeviceWrapper):
+    """Wrapper for managing climate swing mode operations across multiple DPCodes."""
+
+    on_off: DPCodeBooleanWrapper | None = None
+    horizontal: DPCodeBooleanWrapper | None = None
+    vertical: DPCodeBooleanWrapper | None = None
+    modes: list[str]
+
+    @classmethod
+    def find_dpcode(cls, device: CustomerDevice) -> Self | None:
+        """Find and return a _SwingModeWrapper for the given DP codes."""
+        on_off = DPCodeBooleanWrapper.find_dpcode(
+            device, (DPCode.SWING, DPCode.SHAKE), prefer_function=True
+        )
+        horizontal = DPCodeBooleanWrapper.find_dpcode(
+            device, DPCode.SWITCH_HORIZONTAL, prefer_function=True
+        )
+        vertical = DPCodeBooleanWrapper.find_dpcode(
+            device, DPCode.SWITCH_VERTICAL, prefer_function=True
+        )
+        if on_off or horizontal or vertical:
+            modes = [SWING_OFF]
+            if on_off:
+                modes.append(SWING_ON)
+            if horizontal:
+                modes.append(SWING_HORIZONTAL)
+            if vertical:
+                modes.append(SWING_VERTICAL)
+            return cls(
+                on_off=on_off,
+                horizontal=horizontal,
+                vertical=vertical,
+                modes=modes,
+            )
+        return None
+
+    def read_device_status(self, device: CustomerDevice) -> str | None:
+        """Read the device swing mode."""
+        if self.on_off and self.on_off.read_device_status(device):
+            return SWING_ON
+
+        horizontal = (
+            self.horizontal.read_device_status(device) if self.horizontal else None
+        )
+        vertical = self.vertical.read_device_status(device) if self.vertical else None
+        if horizontal and vertical:
+            return SWING_BOTH
+        if horizontal:
+            return SWING_HORIZONTAL
+        if vertical:
+            return SWING_VERTICAL
+
+        return SWING_OFF
+
+    def get_update_commands(
+        self, device: CustomerDevice, value: str
+    ) -> list[dict[str, Any]]:
+        """Set new target swing operation."""
+        commands = []
+        if self.on_off:
+            commands.extend(self.on_off.get_update_commands(device, value == SWING_ON))
+
+        if self.vertical:
+            commands.extend(
+                self.vertical.get_update_commands(
+                    device, value in (SWING_BOTH, SWING_VERTICAL)
+                )
+            )
+        if self.horizontal:
+            commands.extend(
+                self.horizontal.get_update_commands(
+                    device, value in (SWING_BOTH, SWING_HORIZONTAL)
+                )
+            )
+        return commands
+
+
 @dataclass(frozen=True, kw_only=True)
 class TuyaClimateEntityDescription(ClimateEntityDescription):
     """Describe an Tuya climate entity."""
@@ -205,15 +288,7 @@ def async_discover_device(device_ids: list[str]) -> None:
                             device, DPCode.MODE, prefer_function=True
                         ),
                         set_temperature_wrapper=temperature_wrappers[1],
-                        swing_wrapper=DPCodeBooleanWrapper.find_dpcode(
-                            device, (DPCode.SWING, DPCode.SHAKE), prefer_function=True
-                        ),
-                        swing_h_wrapper=DPCodeBooleanWrapper.find_dpcode(
-                            device, DPCode.SWITCH_HORIZONTAL, prefer_function=True
-                        ),
-                        swing_v_wrapper=DPCodeBooleanWrapper.find_dpcode(
-                            device, DPCode.SWITCH_VERTICAL, prefer_function=True
-                        ),
+                        swing_wrapper=_SwingModeWrapper.find_dpcode(device),
                         switch_wrapper=DPCodeBooleanWrapper.find_dpcode(
                             device, DPCode.SWITCH, prefer_function=True
                         ),
@@ -250,9 +325,7 @@ def __init__(
         fan_mode_wrapper: DPCodeEnumWrapper | None,
         hvac_mode_wrapper: DPCodeEnumWrapper | None,
         set_temperature_wrapper: DPCodeIntegerWrapper | None,
-        swing_wrapper: DPCodeBooleanWrapper | None,
-        swing_h_wrapper: DPCodeBooleanWrapper | None,
-        swing_v_wrapper: DPCodeBooleanWrapper | None,
+        swing_wrapper: _SwingModeWrapper | None,
         switch_wrapper: DPCodeBooleanWrapper | None,
         target_humidity_wrapper: _RoundedIntegerWrapper | None,
         temperature_unit: UnitOfTemperature,
@@ -268,8 +341,6 @@ def __init__(
         self._hvac_mode_wrapper = hvac_mode_wrapper
         self._set_temperature = set_temperature_wrapper
         self._swing_wrapper = swing_wrapper
-        self._swing_h_wrapper = swing_h_wrapper
-        self._swing_v_wrapper = swing_v_wrapper
         self._switch_wrapper = switch_wrapper
         self._target_humidity_wrapper = target_humidity_wrapper
         self._attr_temperature_unit = temperature_unit
@@ -324,17 +395,9 @@ def __init__(
             self._attr_fan_modes = fan_mode_wrapper.type_information.range
 
         # Determine swing modes
-        if swing_wrapper or swing_h_wrapper or swing_v_wrapper:
+        if swing_wrapper:
             self._attr_supported_features |= ClimateEntityFeature.SWING_MODE
-            self._attr_swing_modes = [SWING_OFF]
-            if swing_wrapper:
-                self._attr_swing_modes.append(SWING_ON)
-
-            if swing_h_wrapper:
-                self._attr_swing_modes.append(SWING_HORIZONTAL)
-
-            if swing_v_wrapper:
-                self._attr_swing_modes.append(SWING_VERTICAL)
+            self._attr_swing_modes = swing_wrapper.modes
 
         if switch_wrapper:
             self._attr_supported_features |= (
@@ -372,27 +435,7 @@ async def async_set_humidity(self, humidity: int) -> None:
 
     async def async_set_swing_mode(self, swing_mode: str) -> None:
         """Set new target swing operation."""
-        commands = []
-        if self._swing_wrapper:
-            commands.extend(
-                self._swing_wrapper.get_update_commands(
-                    self.device, swing_mode == SWING_ON
-                )
-            )
-        if self._swing_v_wrapper:
-            commands.extend(
-                self._swing_v_wrapper.get_update_commands(
-                    self.device, swing_mode in (SWING_BOTH, SWING_VERTICAL)
-                )
-            )
-        if self._swing_h_wrapper:
-            commands.extend(
-                self._swing_h_wrapper.get_update_commands(
-                    self.device, swing_mode in (SWING_BOTH, SWING_HORIZONTAL)
-                )
-            )
-        if commands:
-            await self._async_send_commands(commands)
+        await self._async_send_wrapper_updates(self._swing_wrapper, swing_mode)
 
     async def async_set_temperature(self, **kwargs: Any) -> None:
         """Set new target temperature."""
@@ -457,21 +500,9 @@ def fan_mode(self) -> str | None:
         return self._read_wrapper(self._fan_mode_wrapper)
 
     @property
-    def swing_mode(self) -> str:
+    def swing_mode(self) -> str | None:
         """Return swing mode."""
-        if self._read_wrapper(self._swing_wrapper):
-            return SWING_ON
-
-        horizontal = self._read_wrapper(self._swing_h_wrapper)
-        vertical = self._read_wrapper(self._swing_v_wrapper)
-        if horizontal and vertical:
-            return SWING_BOTH
-        if horizontal:
-            return SWING_HORIZONTAL
-        if vertical:
-            return SWING_VERTICAL
-
-        return SWING_OFF
+        return self._read_wrapper(self._swing_wrapper)
 
     async def async_turn_on(self) -> None:
         """Turn the device on, retaining current HVAC (if supported)."""

homeassistant/components/tuya/entity.py

@@ -11,7 +11,7 @@
 from homeassistant.helpers.entity import Entity
 
 from .const import DOMAIN, LOGGER, TUYA_HA_SIGNAL_UPDATE_ENTITY
-from .models import DPCodeWrapper
+from .models import DeviceWrapper
 
 
 class TuyaEntity(Entity):
@@ -64,18 +64,20 @@ async def _handle_state_update(
     async def _async_send_commands(self, commands: list[dict[str, Any]]) -> None:
         """Send a list of commands to the device."""
         LOGGER.debug("Sending commands for device %s: %s", self.device.id, commands)
+        if not commands:
+            return
         await self.hass.async_add_executor_job(
             self.device_manager.send_commands, self.device.id, commands
         )
 
-    def _read_wrapper(self, wrapper: DPCodeWrapper | None) -> Any | None:
+    def _read_wrapper(self, wrapper: DeviceWrapper | None) -> Any | None:
         """Read the wrapper device status."""
         if wrapper is None:
             return None
         return wrapper.read_device_status(self.device)
 
     async def _async_send_wrapper_updates(
-        self, wrapper: DPCodeWrapper | None, value: Any
+        self, wrapper: DeviceWrapper | None, value: Any
     ) -> None:
         """Send command to the device."""
         if wrapper is None:

homeassistant/components/tuya/models.py

@@ -165,8 +165,22 @@ def from_json(cls, dpcode: str, type_data: str) -> Self | None:
 }
 
 
-class DPCodeWrapper:
-    """Base DPCode wrapper.
+class DeviceWrapper:
+    """Base device wrapper."""
+
+    def read_device_status(self, device: CustomerDevice) -> Any | None:
+        """Read device status and convert to a Home Assistant value."""
+        raise NotImplementedError
+
+    def get_update_commands(
+        self, device: CustomerDevice, value: Any
+    ) -> list[dict[str, Any]]:
+        """Generate update commands for a Home Assistant action."""
+        raise NotImplementedError
+
+
+class DPCodeWrapper(DeviceWrapper):
+    """Base device wrapper for a single DPCode.
 
     Used as a common interface for referring to a DPCode, and
     access read conversion routines.
@@ -186,13 +200,6 @@ def _read_device_status_raw(self, device: CustomerDevice) -> Any | None:
         """
         return device.status.get(self.dpcode)
 
-    def read_device_status(self, device: CustomerDevice) -> Any | None:
-        """Read the device value for the dpcode.
-
-        The raw device status is converted to a Home Assistant value.
-        """
-        raise NotImplementedError
-
     def _convert_value_to_raw_value(self, device: CustomerDevice, value: Any) -> Any:
         """Convert a Home Assistant value back to a raw device value.