core/homeassistant/components/tuya/valve.py

"""Support for Tuya valves."""

from __future__ import annotations

from tuya_device_handlers.definition.valve import (
    TuyaValveDefinition,
    get_default_definition,
)
from tuya_sharing import CustomerDevice, Manager

from homeassistant.components.valve import (
    ValveDeviceClass,
    ValveEntity,
    ValveEntityDescription,
    ValveEntityFeature,
)
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback

from . import TuyaConfigEntry
from .const import TUYA_DISCOVERY_NEW, DeviceCategory, DPCode
from .entity import TuyaEntity

VALVES: dict[DeviceCategory, tuple[ValveEntityDescription, ...]] = {
    DeviceCategory.SFKZQ: (
        ValveEntityDescription(
            key=DPCode.SWITCH,
            translation_key="valve",
            device_class=ValveDeviceClass.WATER,
        ),
        ValveEntityDescription(
            key=DPCode.SWITCH_1,
            translation_key="indexed_valve",
            translation_placeholders={"index": "1"},
            device_class=ValveDeviceClass.WATER,
        ),
        ValveEntityDescription(
            key=DPCode.SWITCH_2,
            translation_key="indexed_valve",
            translation_placeholders={"index": "2"},
            device_class=ValveDeviceClass.WATER,
        ),
        ValveEntityDescription(
            key=DPCode.SWITCH_3,
            translation_key="indexed_valve",
            translation_placeholders={"index": "3"},
            device_class=ValveDeviceClass.WATER,
        ),
        ValveEntityDescription(
            key=DPCode.SWITCH_4,
            translation_key="indexed_valve",
            translation_placeholders={"index": "4"},
            device_class=ValveDeviceClass.WATER,
        ),
        ValveEntityDescription(
            key=DPCode.SWITCH_5,
            translation_key="indexed_valve",
            translation_placeholders={"index": "5"},
            device_class=ValveDeviceClass.WATER,
        ),
        ValveEntityDescription(
            key=DPCode.SWITCH_6,
            translation_key="indexed_valve",
            translation_placeholders={"index": "6"},
            device_class=ValveDeviceClass.WATER,
        ),
        ValveEntityDescription(
            key=DPCode.SWITCH_7,
            translation_key="indexed_valve",
            translation_placeholders={"index": "7"},
            device_class=ValveDeviceClass.WATER,
        ),
        ValveEntityDescription(
            key=DPCode.SWITCH_8,
            translation_key="indexed_valve",
            translation_placeholders={"index": "8"},
            device_class=ValveDeviceClass.WATER,
        ),
    ),
}


async def async_setup_entry(
    hass: HomeAssistant,
    entry: TuyaConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up tuya valves dynamically through tuya discovery."""
    manager = entry.runtime_data.manager

    @callback
    def async_discover_device(device_ids: list[str]) -> None:
        """Discover and add a discovered tuya valve."""
        entities: list[TuyaValveEntity] = []
        for device_id in device_ids:
            device = manager.device_map[device_id]
            if descriptions := VALVES.get(device.category):
                entities.extend(
                    TuyaValveEntity(device, manager, description, definition)
                    for description in descriptions
                    if (definition := get_default_definition(device, description.key))
                )

        async_add_entities(entities)

    async_discover_device([*manager.device_map])

    entry.async_on_unload(
        async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device)
    )


class TuyaValveEntity(TuyaEntity, ValveEntity):
    """Tuya Valve Device."""

    _attr_supported_features = ValveEntityFeature.OPEN | ValveEntityFeature.CLOSE

    def __init__(
        self,
        device: CustomerDevice,
        device_manager: Manager,
        description: ValveEntityDescription,
        definition: TuyaValveDefinition,
    ) -> None:
        """Init TuyaValveEntity."""
        super().__init__(device, device_manager, description)
        self._dpcode_wrapper = definition.control_wrapper

    @property
    def is_closed(self) -> bool | None:
        """Return if the valve is closed."""
        if (is_open := self._read_wrapper(self._dpcode_wrapper)) is None:
            return None
        return not is_open

    async def _process_device_update(
        self,
        updated_status_properties: list[str],
        dp_timestamps: dict[str, int] | None,
    ) -> bool:
        """Called when Tuya device sends an update with updated properties.

        Returns True if the Home Assistant state should be written,
        or False if the state write should be skipped.
        """
        return not self._dpcode_wrapper.skip_update(
            self.device, updated_status_properties, dp_timestamps
        )

    async def async_open_valve(self) -> None:
        """Open the valve."""
        await self._async_send_wrapper_updates(self._dpcode_wrapper, True)

    async def async_close_valve(self) -> None:
        """Close the valve."""
        await self._async_send_wrapper_updates(self._dpcode_wrapper, False)