fan.py

"""Support for Rabbit Air fan entity."""

from typing import Any

from rabbitair import Mode, Model, Speed

from homeassistant.components.fan import FanEntity, FanEntityFeature
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
from homeassistant.util.percentage import (
    ordered_list_item_to_percentage,
    percentage_to_ordered_list_item,
)

from .coordinator import RabbitAirConfigEntry, RabbitAirDataUpdateCoordinator
from .entity import RabbitAirBaseEntity

SPEED_LIST = [
    Speed.Silent,
    Speed.Low,
    Speed.Medium,
    Speed.High,
    Speed.Turbo,
]

PRESET_MODE_AUTO = "auto"
PRESET_MODE_MANUAL = "manual"
PRESET_MODE_POLLEN = "pollen"

PRESET_MODES = {
    PRESET_MODE_AUTO: Mode.Auto,
    PRESET_MODE_MANUAL: Mode.Manual,
    PRESET_MODE_POLLEN: Mode.Pollen,
}


async def async_setup_entry(
    hass: HomeAssistant,
    entry: RabbitAirConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up a config entry."""
    async_add_entities([RabbitAirFanEntity(entry.runtime_data, entry)])


class RabbitAirFanEntity(RabbitAirBaseEntity, FanEntity):
    """Fan control functions of the Rabbit Air air purifier."""

    _attr_translation_key = "rabbitair"
    _attr_supported_features = (
        FanEntityFeature.PRESET_MODE
        | FanEntityFeature.SET_SPEED
        | FanEntityFeature.TURN_ON
        | FanEntityFeature.TURN_OFF
    )

    def __init__(
        self,
        coordinator: RabbitAirDataUpdateCoordinator,
        entry: RabbitAirConfigEntry,
    ) -> None:
        """Initialize the entity."""
        super().__init__(coordinator, entry)
        self._attr_name = entry.title

        if self._is_model(Model.MinusA2):
            self._attr_preset_modes = list(PRESET_MODES)
        elif self._is_model(Model.A3):
            # A3 does not support Pollen mode
            self._attr_preset_modes = [
                k for k in PRESET_MODES if k != PRESET_MODE_POLLEN
            ]

        self._attr_speed_count = len(SPEED_LIST)

        self._get_state_from_coordinator_data()

    @callback
    def _handle_coordinator_update(self) -> None:
        """Handle updated data from the coordinator."""
        self._get_state_from_coordinator_data()
        super()._handle_coordinator_update()

    def _get_state_from_coordinator_data(self) -> None:
        """Populate the entity fields with values from the coordinator data."""
        data = self.coordinator.data

        # Speed as a percentage
        if not data.power:
            self._attr_percentage = 0
        elif data.speed is None:
            self._attr_percentage = None
        elif data.speed is Speed.SuperSilent:
            self._attr_percentage = 1
        else:
            self._attr_percentage = ordered_list_item_to_percentage(
                SPEED_LIST, data.speed
            )

        # Preset mode
        if not data.power or data.mode is None:
            self._attr_preset_mode = None
        else:
            # Get key by value in dictionary
            self._attr_preset_mode = next(
                k for k, v in PRESET_MODES.items() if v is data.mode
            )

    async def async_set_preset_mode(self, preset_mode: str) -> None:
        """Set new preset mode."""
        await self._set_state(power=True, mode=PRESET_MODES[preset_mode])
        self._attr_preset_mode = preset_mode
        self.async_write_ha_state()

    async def async_set_percentage(self, percentage: int) -> None:
        """Set the speed of the fan, as a percentage."""
        if percentage > 0:
            value = percentage_to_ordered_list_item(SPEED_LIST, percentage)
            await self._set_state(power=True, speed=value)
            self._attr_percentage = percentage
        else:
            await self._set_state(power=False)
            self._attr_percentage = 0
            self._attr_preset_mode = None
        self.async_write_ha_state()

    async def async_turn_on(
        self,
        percentage: int | None = None,
        preset_mode: str | None = None,
        **kwargs: Any,
    ) -> None:
        """Turn on the fan."""
        mode_value: Mode | None = None
        if preset_mode is not None:
            mode_value = PRESET_MODES[preset_mode]
        speed_value: Speed | None = None
        if percentage is not None:
            speed_value = percentage_to_ordered_list_item(SPEED_LIST, percentage)
        await self._set_state(power=True, mode=mode_value, speed=speed_value)
        if percentage is not None:
            self._attr_percentage = percentage
        if preset_mode is not None:
            self._attr_preset_mode = preset_mode
        self.async_write_ha_state()

    async def async_turn_off(self, **kwargs: Any) -> None:
        """Turn the fan off."""
        await self._set_state(power=False)
        self._attr_percentage = 0
        self._attr_preset_mode = None
        self.async_write_ha_state()