hass-core/homeassistant/components/rainbird/switch.py

"""Support for Rain Bird Irrigation system LNK WiFi Module."""

import logging

from pyrainbird import AvailableStations, RainbirdController
import voluptuous as vol

from homeassistant.components.switch import SwitchDevice
from homeassistant.const import ATTR_ENTITY_ID, CONF_FRIENDLY_NAME, CONF_TRIGGER_TIME
from homeassistant.helpers import config_validation as cv

from . import CONF_ZONES, DATA_RAINBIRD, DOMAIN, RAINBIRD_CONTROLLER

_LOGGER = logging.getLogger(__name__)

ATTR_DURATION = "duration"

SERVICE_START_IRRIGATION = "start_irrigation"

SERVICE_SCHEMA_IRRIGATION = vol.Schema(
    {
        vol.Required(ATTR_ENTITY_ID): cv.entity_id,
        vol.Required(ATTR_DURATION): vol.All(vol.Coerce(float), vol.Range(min=0)),
    }
)


def setup_platform(hass, config, add_entities, discovery_info=None):
    """Set up Rain Bird switches over a Rain Bird controller."""

    if discovery_info is None:
        return

    controller: RainbirdController = hass.data[DATA_RAINBIRD][
        discovery_info[RAINBIRD_CONTROLLER]
    ]
    available_stations: AvailableStations = controller.get_available_stations()
    if not (available_stations and available_stations.stations):
        return
    devices = []
    for zone in range(1, available_stations.stations.count + 1):
        if available_stations.stations.active(zone):
            zone_config = discovery_info.get(CONF_ZONES, {}).get(zone, {})
            time = zone_config.get(CONF_TRIGGER_TIME, discovery_info[CONF_TRIGGER_TIME])
            name = zone_config.get(CONF_FRIENDLY_NAME)
            devices.append(
                RainBirdSwitch(
                    controller,
                    zone,
                    time,
                    name if name else "Sprinkler {}".format(zone),
                )
            )

    add_entities(devices, True)

    def start_irrigation(service):
        entity_id = service.data[ATTR_ENTITY_ID]
        duration = service.data[ATTR_DURATION]

        for device in devices:
            if device.entity_id == entity_id:
                device.turn_on(duration=duration)

    hass.services.register(
        DOMAIN,
        SERVICE_START_IRRIGATION,
        start_irrigation,
        schema=SERVICE_SCHEMA_IRRIGATION,
    )


class RainBirdSwitch(SwitchDevice):
    """Representation of a Rain Bird switch."""

    def __init__(self, controller: RainbirdController, zone, time, name):
        """Initialize a Rain Bird Switch Device."""
        self._rainbird = controller
        self._zone = zone
        self._name = name
        self._state = None
        self._duration = time
        self._attributes = {ATTR_DURATION: self._duration, "zone": self._zone}

    @property
    def device_state_attributes(self):
        """Return state attributes."""
        return self._attributes

    @property
    def name(self):
        """Get the name of the switch."""
        return self._name

    def turn_on(self, **kwargs):
        """Turn the switch on."""
        if self._rainbird.irrigate_zone(
            int(self._zone),
            int(kwargs[ATTR_DURATION] if ATTR_DURATION in kwargs else self._duration),
        ):
            self._state = True

    def turn_off(self, **kwargs):
        """Turn the switch off."""
        if self._rainbird.stop_irrigation():
            self._state = False

    def update(self):
        """Update switch status."""
        self._state = self._rainbird.get_zone_state(self._zone)

    @property
    def is_on(self):
        """Return true if switch is on."""
        return self._state