core/homeassistant/components/esphome/__init__.py

"""Support for esphome devices."""
import asyncio
import logging
import math
from typing import Any, Dict, List, Optional, TYPE_CHECKING, Callable, Tuple

import attr
import voluptuous as vol

from homeassistant import const
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, \
    EVENT_HOMEASSISTANT_STOP
from homeassistant.core import callback, Event, State
import homeassistant.helpers.device_registry as dr
from homeassistant.exceptions import TemplateError
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers import template
from homeassistant.helpers.dispatcher import async_dispatcher_connect, \
    async_dispatcher_send
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.event import async_track_state_change
from homeassistant.helpers.template import Template
from homeassistant.helpers.json import JSONEncoder
from homeassistant.helpers.storage import Store
from homeassistant.helpers.typing import HomeAssistantType, ConfigType

# Import config flow so that it's added to the registry
from .config_flow import EsphomeFlowHandler  # noqa

if TYPE_CHECKING:
    from aioesphomeapi import APIClient, EntityInfo, EntityState, DeviceInfo, \
        ServiceCall, UserService

DOMAIN = 'esphome'
_LOGGER = logging.getLogger(__name__)

DISPATCHER_UPDATE_ENTITY = 'esphome_{entry_id}_update_{component_key}_{key}'
DISPATCHER_REMOVE_ENTITY = 'esphome_{entry_id}_remove_{component_key}_{key}'
DISPATCHER_ON_LIST = 'esphome_{entry_id}_on_list'
DISPATCHER_ON_DEVICE_UPDATE = 'esphome_{entry_id}_on_device_update'
DISPATCHER_ON_STATE = 'esphome_{entry_id}_on_state'

STORAGE_KEY = 'esphome.{}'
STORAGE_VERSION = 1

# The HA component types this integration supports
HA_COMPONENTS = [
    'binary_sensor',
    'camera',
    'climate',
    'cover',
    'fan',
    'light',
    'sensor',
    'switch',
]

# No config schema - only configuration entry
CONFIG_SCHEMA = vol.Schema({}, extra=vol.ALLOW_EXTRA)


@attr.s
class RuntimeEntryData:
    """Store runtime data for esphome config entries."""

    entry_id = attr.ib(type=str)
    client = attr.ib(type='APIClient')
    store = attr.ib(type=Store)
    reconnect_task = attr.ib(type=Optional[asyncio.Task], default=None)
    state = attr.ib(type=Dict[str, Dict[str, Any]], factory=dict)
    info = attr.ib(type=Dict[str, Dict[str, Any]], factory=dict)
    services = attr.ib(type=Dict[int, 'UserService'], factory=dict)
    available = attr.ib(type=bool, default=False)
    device_info = attr.ib(type='DeviceInfo', default=None)
    cleanup_callbacks = attr.ib(type=List[Callable[[], None]], factory=list)
    disconnect_callbacks = attr.ib(type=List[Callable[[], None]], factory=list)

    def async_update_entity(self, hass: HomeAssistantType, component_key: str,
                            key: int) -> None:
        """Schedule the update of an entity."""
        signal = DISPATCHER_UPDATE_ENTITY.format(
            entry_id=self.entry_id, component_key=component_key, key=key)
        async_dispatcher_send(hass, signal)

    def async_remove_entity(self, hass: HomeAssistantType, component_key: str,
                            key: int) -> None:
        """Schedule the removal of an entity."""
        signal = DISPATCHER_REMOVE_ENTITY.format(
            entry_id=self.entry_id, component_key=component_key, key=key)
        async_dispatcher_send(hass, signal)

    def async_update_static_infos(self, hass: HomeAssistantType,
                                  infos: 'List[EntityInfo]') -> None:
        """Distribute an update of static infos to all platforms."""
        signal = DISPATCHER_ON_LIST.format(entry_id=self.entry_id)
        async_dispatcher_send(hass, signal, infos)

    def async_update_state(self, hass: HomeAssistantType,
                           state: 'EntityState') -> None:
        """Distribute an update of state information to all platforms."""
        signal = DISPATCHER_ON_STATE.format(entry_id=self.entry_id)
        async_dispatcher_send(hass, signal, state)

    def async_update_device_state(self, hass: HomeAssistantType) -> None:
        """Distribute an update of a core device state like availability."""
        signal = DISPATCHER_ON_DEVICE_UPDATE.format(entry_id=self.entry_id)
        async_dispatcher_send(hass, signal)

    async def async_load_from_store(self) -> Tuple[List['EntityInfo'],
                                                   List['UserService']]:
        """Load the retained data from store and return de-serialized data."""
        # pylint: disable= redefined-outer-name
        from aioesphomeapi import COMPONENT_TYPE_TO_INFO, DeviceInfo, \
            UserService

        restored = await self.store.async_load()
        if restored is None:
            return [], []

        self.device_info = _attr_obj_from_dict(DeviceInfo,
                                               **restored.pop('device_info'))
        infos = []
        for comp_type, restored_infos in restored.items():
            if comp_type not in COMPONENT_TYPE_TO_INFO:
                continue
            for info in restored_infos:
                cls = COMPONENT_TYPE_TO_INFO[comp_type]
                infos.append(_attr_obj_from_dict(cls, **info))
        services = []
        for service in restored.get('services', []):
            services.append(UserService.from_dict(service))
        return infos, services

    async def async_save_to_store(self) -> None:
        """Generate dynamic data to store and save it to the filesystem."""
        store_data = {
            'device_info': attr.asdict(self.device_info),
            'services': []
        }

        for comp_type, infos in self.info.items():
            store_data[comp_type] = [attr.asdict(info)
                                     for info in infos.values()]
        for service in self.services.values():
            store_data['services'].append(service.to_dict())

        await self.store.async_save(store_data)


def _attr_obj_from_dict(cls, **kwargs):
    return cls(**{key: kwargs[key] for key in attr.fields_dict(cls)
                  if key in kwargs})


async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
    """Stub to allow setting up this component.

    Configuration through YAML is not supported at this time.
    """
    return True


async def async_setup_entry(hass: HomeAssistantType,
                            entry: ConfigEntry) -> bool:
    """Set up the esphome component."""
    # pylint: disable=redefined-outer-name
    from aioesphomeapi import APIClient, APIConnectionError

    hass.data.setdefault(DOMAIN, {})

    host = entry.data[CONF_HOST]
    port = entry.data[CONF_PORT]
    password = entry.data[CONF_PASSWORD]

    cli = APIClient(hass.loop, host, port, password,
                    client_info="Home Assistant {}".format(const.__version__))

    # Store client in per-config-entry hass.data
    store = Store(hass, STORAGE_VERSION, STORAGE_KEY.format(entry.entry_id),
                  encoder=JSONEncoder)
    entry_data = hass.data[DOMAIN][entry.entry_id] = RuntimeEntryData(
        client=cli,
        entry_id=entry.entry_id,
        store=store,
    )

    async def on_stop(event: Event) -> None:
        """Cleanup the socket client on HA stop."""
        await _cleanup_instance(hass, entry)

    entry_data.cleanup_callbacks.append(
        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, on_stop)
    )

    @callback
    def async_on_state(state: 'EntityState') -> None:
        """Send dispatcher updates when a new state is received."""
        entry_data.async_update_state(hass, state)

    @callback
    def async_on_service_call(service: 'ServiceCall') -> None:
        """Call service when user automation in ESPHome config is triggered."""
        domain, service_name = service.service.split('.', 1)
        service_data = service.data

        if service.data_template:
            try:
                data_template = {key: Template(value) for key, value in
                                 service.data_template.items()}
                template.attach(hass, data_template)
                service_data.update(template.render_complex(
                    data_template, service.variables))
            except TemplateError as ex:
                _LOGGER.error('Error rendering data template: %s', ex)
                return

        hass.async_create_task(hass.services.async_call(
            domain, service_name, service_data, blocking=True))

    async def send_home_assistant_state(entity_id: str, _,
                                        new_state: Optional[State]) -> None:
        """Forward Home Assistant states to ESPHome."""
        if new_state is None:
            return
        await cli.send_home_assistant_state(entity_id, new_state.state)

    @callback
    def async_on_state_subscription(entity_id: str) -> None:
        """Subscribe and forward states for requested entities."""
        unsub = async_track_state_change(
            hass, entity_id, send_home_assistant_state)
        entry_data.disconnect_callbacks.append(unsub)
        # Send initial state
        hass.async_create_task(send_home_assistant_state(
            entity_id, None, hass.states.get(entity_id)))

    async def on_login() -> None:
        """Subscribe to states and list entities on successful API login."""
        try:
            entry_data.device_info = await cli.device_info()
            entry_data.available = True
            await _async_setup_device_registry(hass, entry,
                                               entry_data.device_info)
            entry_data.async_update_device_state(hass)

            entity_infos, services = await cli.list_entities_services()
            entry_data.async_update_static_infos(hass, entity_infos)
            await _setup_services(hass, entry_data, services)
            await cli.subscribe_states(async_on_state)
            await cli.subscribe_service_calls(async_on_service_call)
            await cli.subscribe_home_assistant_states(
                async_on_state_subscription)

            hass.async_create_task(entry_data.async_save_to_store())
        except APIConnectionError as err:
            _LOGGER.warning("Error getting initial data: %s", err)
            # Re-connection logic will trigger after this
            await cli.disconnect()

    try_connect = await _setup_auto_reconnect_logic(hass, cli, entry, host,
                                                    on_login)

    # This is a bit of a hack: We schedule complete_setup into the
    # event loop and return immediately (return True)
    #
    # Usually, we should avoid that so that HA can track which components
    # have been started successfully and which failed to be set up.
    # That doesn't work here for two reasons:
    #  - We have our own re-connect logic
    #  - Before we do the first try_connect() call, we need to make sure
    #    all dispatcher event listeners have been connected, so
    #    async_forward_entry_setup needs to be awaited. However, if we
    #    would await async_forward_entry_setup() in async_setup_entry(),
    #    we would end up with a deadlock.
    #
    # Solution is: complete the setup outside of the async_setup_entry()
    # function. HA will wait until the first connection attempt is made
    # before starting up (as it should), but if the first connection attempt
    # fails we will schedule all next re-connect attempts outside of the
    # tracked tasks (hass.loop.create_task). This way HA won't stall startup
    # forever until a connection is successful.

    async def complete_setup() -> None:
        """Complete the config entry setup."""
        tasks = []
        for component in HA_COMPONENTS:
            tasks.append(hass.config_entries.async_forward_entry_setup(
                entry, component))
        await asyncio.wait(tasks)

        infos, services = await entry_data.async_load_from_store()
        entry_data.async_update_static_infos(hass, infos)
        await _setup_services(hass, entry_data, services)

        # If first connect fails, the next re-connect will be scheduled
        # outside of _pending_task, in order not to delay HA startup
        # indefinitely
        await try_connect(is_disconnect=False)

    hass.async_create_task(complete_setup())
    return True


async def _setup_auto_reconnect_logic(hass: HomeAssistantType,
                                      cli: 'APIClient',
                                      entry: ConfigEntry, host: str, on_login):
    """Set up the re-connect logic for the API client."""
    from aioesphomeapi import APIConnectionError

    async def try_connect(tries: int = 0, is_disconnect: bool = True) -> None:
        """Try connecting to the API client. Will retry if not successful."""
        if entry.entry_id not in hass.data[DOMAIN]:
            # When removing/disconnecting manually
            return

        data = hass.data[DOMAIN][entry.entry_id]  # type: RuntimeEntryData
        for disconnect_cb in data.disconnect_callbacks:
            disconnect_cb()
        data.disconnect_callbacks = []
        data.available = False
        data.async_update_device_state(hass)

        if is_disconnect:
            # This can happen often depending on WiFi signal strength.
            # So therefore all these connection warnings are logged
            # as infos. The "unavailable" logic will still trigger so the
            # user knows if the device is not connected.
            _LOGGER.info("Disconnected from ESPHome API for %s", host)

        if tries != 0:
            # If not first re-try, wait and print message
            # Cap wait time at 1 minute. This is because while working on the
            # device (e.g. soldering stuff), users don't want to have to wait
            # a long time for their device to show up in HA again (this was
            # mentioned a lot in early feedback)
            #
            # In the future another API will be set up so that the ESP can
            # notify HA of connectivity directly, but for new we'll use a
            # really short reconnect interval.
            tries = min(tries, 10)  # prevent OverflowError
            wait_time = int(round(min(1.8**tries, 60.0)))
            _LOGGER.info("Trying to reconnect in %s seconds", wait_time)
            await asyncio.sleep(wait_time)

        try:
            await cli.connect(on_stop=try_connect, login=True)
        except APIConnectionError as error:
            _LOGGER.info("Can't connect to ESPHome API for %s: %s",
                         host, error)
            # Schedule re-connect in event loop in order not to delay HA
            # startup. First connect is scheduled in tracked tasks.
            data.reconnect_task = hass.loop.create_task(
                try_connect(tries + 1, is_disconnect=False))
        else:
            _LOGGER.info("Successfully connected to %s", host)
            hass.async_create_task(on_login())

    return try_connect


async def _async_setup_device_registry(hass: HomeAssistantType,
                                       entry: ConfigEntry,
                                       device_info: 'DeviceInfo'):
    """Set up device registry feature for a particular config entry."""
    sw_version = device_info.esphome_core_version
    if device_info.compilation_time:
        sw_version += ' ({})'.format(device_info.compilation_time)
    device_registry = await dr.async_get_registry(hass)
    device_registry.async_get_or_create(
        config_entry_id=entry.entry_id,
        connections={
            (dr.CONNECTION_NETWORK_MAC, device_info.mac_address)
        },
        name=device_info.name,
        manufacturer='espressif',
        model=device_info.model,
        sw_version=sw_version,
    )


async def _register_service(hass: HomeAssistantType,
                            entry_data: RuntimeEntryData,
                            service: 'UserService'):
    from aioesphomeapi import UserServiceArgType
    service_name = '{}_{}'.format(entry_data.device_info.name, service.name)
    schema = {}
    for arg in service.args:
        schema[vol.Required(arg.name)] = {
            UserServiceArgType.BOOL: cv.boolean,
            UserServiceArgType.INT: vol.Coerce(int),
            UserServiceArgType.FLOAT: vol.Coerce(float),
            UserServiceArgType.STRING: cv.string,
        }[arg.type_]

    async def execute_service(call):
        await entry_data.client.execute_service(service, call.data)

    hass.services.async_register(DOMAIN, service_name, execute_service,
                                 vol.Schema(schema))


async def _setup_services(hass: HomeAssistantType,
                          entry_data: RuntimeEntryData,
                          services: List['UserService']):
    old_services = entry_data.services.copy()
    to_unregister = []
    to_register = []
    for service in services:
        if service.key in old_services:
            # Already exists
            matching = old_services.pop(service.key)
            if matching != service:
                # Need to re-register
                to_unregister.append(matching)
                to_register.append(service)
        else:
            # New service
            to_register.append(service)

    for service in old_services.values():
        to_unregister.append(service)

    entry_data.services = {serv.key: serv for serv in services}

    for service in to_unregister:
        service_name = '{}_{}'.format(entry_data.device_info.name,
                                      service.name)
        hass.services.async_remove(DOMAIN, service_name)

    for service in to_register:
        await _register_service(hass, entry_data, service)


async def _cleanup_instance(hass: HomeAssistantType,
                            entry: ConfigEntry) -> None:
    """Cleanup the esphome client if it exists."""
    data = hass.data[DOMAIN].pop(entry.entry_id)  # type: RuntimeEntryData
    if data.reconnect_task is not None:
        data.reconnect_task.cancel()
    for disconnect_cb in data.disconnect_callbacks:
        disconnect_cb()
    for cleanup_callback in data.cleanup_callbacks:
        cleanup_callback()
    await data.client.disconnect()


async def async_unload_entry(hass: HomeAssistantType,
                             entry: ConfigEntry) -> bool:
    """Unload an esphome config entry."""
    await _cleanup_instance(hass, entry)

    tasks = []
    for component in HA_COMPONENTS:
        tasks.append(hass.config_entries.async_forward_entry_unload(
            entry, component))
    await asyncio.wait(tasks)

    return True


async def platform_async_setup_entry(hass: HomeAssistantType,
                                     entry: ConfigEntry,
                                     async_add_entities,
                                     *,
                                     component_key: str,
                                     info_type,
                                     entity_type,
                                     state_type
                                     ) -> None:
    """Set up an esphome platform.

    This method is in charge of receiving, distributing and storing
    info and state updates.
    """
    entry_data = hass.data[DOMAIN][entry.entry_id]  # type: RuntimeEntryData
    entry_data.info[component_key] = {}
    entry_data.state[component_key] = {}

    @callback
    def async_list_entities(infos: List['EntityInfo']):
        """Update entities of this platform when entities are listed."""
        old_infos = entry_data.info[component_key]
        new_infos = {}
        add_entities = []
        for info in infos:
            if not isinstance(info, info_type):
                # Filter out infos that don't belong to this platform.
                continue

            if info.key in old_infos:
                # Update existing entity
                old_infos.pop(info.key)
            else:
                # Create new entity
                entity = entity_type(entry.entry_id, component_key, info.key)
                add_entities.append(entity)
            new_infos[info.key] = info

        # Remove old entities
        for info in old_infos.values():
            entry_data.async_remove_entity(hass, component_key, info.key)
        entry_data.info[component_key] = new_infos
        async_add_entities(add_entities)

    signal = DISPATCHER_ON_LIST.format(entry_id=entry.entry_id)
    entry_data.cleanup_callbacks.append(
        async_dispatcher_connect(hass, signal, async_list_entities)
    )

    @callback
    def async_entity_state(state: 'EntityState'):
        """Notify the appropriate entity of an updated state."""
        if not isinstance(state, state_type):
            return
        entry_data.state[component_key][state.key] = state
        entry_data.async_update_entity(hass, component_key, state.key)

    signal = DISPATCHER_ON_STATE.format(entry_id=entry.entry_id)
    entry_data.cleanup_callbacks.append(
        async_dispatcher_connect(hass, signal, async_entity_state)
    )


def esphome_state_property(func):
    """Wrap a state property of an esphome entity.

    This checks if the state object in the entity is set, and
    prevents writing NAN values to the Home Assistant state machine.
    """
    @property
    def _wrapper(self):
        if self._state is None:
            return None
        val = func(self)
        if isinstance(val, float) and math.isnan(val):
            # Home Assistant doesn't use NAN values in state machine
            # (not JSON serializable)
            return None
        return val
    return _wrapper


class EsphomeEnumMapper:
    """Helper class to convert between hass and esphome enum values."""

    def __init__(self, func: Callable[[], Dict[int, str]]):
        """Construct a EsphomeEnumMapper."""
        self._func = func

    def from_esphome(self, value: int) -> str:
        """Convert from an esphome int representation to a hass string."""
        return self._func()[value]

    def from_hass(self, value: str) -> int:
        """Convert from a hass string to a esphome int representation."""
        inverse = {v: k for k, v in self._func().items()}
        return inverse[value]


def esphome_map_enum(func: Callable[[], Dict[int, str]]):
    """Map esphome int enum values to hass string constants.

    This class has to be used as a decorator. This ensures the aioesphomeapi
    import is only happening at runtime.
    """
    return EsphomeEnumMapper(func)


class EsphomeEntity(Entity):
    """Define a generic esphome entity."""

    def __init__(self, entry_id: str, component_key: str, key: int):
        """Initialize."""
        self._entry_id = entry_id
        self._component_key = component_key
        self._key = key
        self._remove_callbacks = []  # type: List[Callable[[], None]]

    async def async_added_to_hass(self) -> None:
        """Register callbacks."""
        kwargs = {
            'entry_id': self._entry_id,
            'component_key': self._component_key,
            'key': self._key,
        }
        self._remove_callbacks.append(
            async_dispatcher_connect(self.hass,
                                     DISPATCHER_UPDATE_ENTITY.format(**kwargs),
                                     self._on_update)
        )

        self._remove_callbacks.append(
            async_dispatcher_connect(self.hass,
                                     DISPATCHER_REMOVE_ENTITY.format(**kwargs),
                                     self.async_remove)
        )

        self._remove_callbacks.append(
            async_dispatcher_connect(
                self.hass, DISPATCHER_ON_DEVICE_UPDATE.format(**kwargs),
                self.async_schedule_update_ha_state)
        )

    async def _on_update(self) -> None:
        """Update the entity state when state or static info changed."""
        self.async_schedule_update_ha_state()

    async def async_will_remove_from_hass(self) -> None:
        """Unregister callbacks."""
        for remove_callback in self._remove_callbacks:
            remove_callback()
        self._remove_callbacks = []

    @property
    def _entry_data(self) -> RuntimeEntryData:
        return self.hass.data[DOMAIN][self._entry_id]

    @property
    def _static_info(self) -> 'EntityInfo':
        return self._entry_data.info[self._component_key][self._key]

    @property
    def _device_info(self) -> 'DeviceInfo':
        return self._entry_data.device_info

    @property
    def _client(self) -> 'APIClient':
        return self._entry_data.client

    @property
    def _state(self) -> 'Optional[EntityState]':
        try:
            return self._entry_data.state[self._component_key][self._key]
        except KeyError:
            return None

    @property
    def available(self) -> bool:
        """Return if the entity is available."""
        device = self._device_info

        if device.has_deep_sleep:
            # During deep sleep the ESP will not be connectable (by design)
            # For these cases, show it as available
            return True

        return self._entry_data.available

    @property
    def unique_id(self) -> Optional[str]:
        """Return a unique id identifying the entity."""
        if not self._static_info.unique_id:
            return None
        return self._static_info.unique_id

    @property
    def device_info(self) -> Dict[str, Any]:
        """Return device registry information for this entity."""
        return {
            'connections': {(dr.CONNECTION_NETWORK_MAC,
                             self._device_info.mac_address)}
        }

    @property
    def name(self) -> str:
        """Return the name of the entity."""
        return self._static_info.name

    @property
    def should_poll(self) -> bool:
        """Disable polling."""
        return False