diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 395c145e5df..338968a4fd7 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -39,18 +39,6 @@ _LOGGER = logging.getLogger(__name__) 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) @@ -144,7 +132,8 @@ async def async_setup_entry(hass: HomeAssistantType, 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 entry_data.async_update_static_infos( + hass, entry, 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) @@ -162,14 +151,8 @@ async def async_setup_entry(hass: HomeAssistantType, 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 entry_data.async_update_static_infos(hass, entry, infos) await _setup_services(hass, entry_data, services) # Create connection attempt outside of HA's tracked task in order @@ -308,7 +291,7 @@ async def _setup_services(hass: HomeAssistantType, async def _cleanup_instance(hass: HomeAssistantType, - entry: ConfigEntry) -> None: + entry: ConfigEntry) -> RuntimeEntryData: """Cleanup the esphome client if it exists.""" data = hass.data[DATA_KEY].pop(entry.entry_id) # type: RuntimeEntryData if data.reconnect_task is not None: @@ -318,19 +301,19 @@ async def _cleanup_instance(hass: HomeAssistantType, for cleanup_callback in data.cleanup_callbacks: cleanup_callback() await data.client.disconnect() + return data async def async_unload_entry(hass: HomeAssistantType, entry: ConfigEntry) -> bool: """Unload an esphome config entry.""" - await _cleanup_instance(hass, entry) - + entry_data = await _cleanup_instance(hass, entry) tasks = [] - for component in HA_COMPONENTS: + for platform in entry_data.loaded_platforms: tasks.append(hass.config_entries.async_forward_entry_unload( - entry, component)) - await asyncio.wait(tasks) - + entry, platform)) + if tasks: + await asyncio.wait(tasks) return True diff --git a/homeassistant/components/esphome/entry_data.py b/homeassistant/components/esphome/entry_data.py index 47cadc00653..4e78718b760 100644 --- a/homeassistant/components/esphome/entry_data.py +++ b/homeassistant/components/esphome/entry_data.py @@ -1,11 +1,15 @@ """Runtime entry data for ESPHome stored in hass.data.""" import asyncio -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple, Set from aioesphomeapi import ( - COMPONENT_TYPE_TO_INFO, DeviceInfo, EntityInfo, EntityState, UserService) + COMPONENT_TYPE_TO_INFO, DeviceInfo, EntityInfo, EntityState, UserService, + BinarySensorInfo, + CameraInfo, ClimateInfo, CoverInfo, FanInfo, LightInfo, SensorInfo, + SwitchInfo, TextSensorInfo) import attr +from homeassistant.config_entries import ConfigEntry from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.storage import Store from homeassistant.helpers.typing import HomeAssistantType @@ -17,6 +21,19 @@ 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' +# Mapping from ESPHome info type to HA platform +INFO_TYPE_TO_PLATFORM = { + BinarySensorInfo: 'binary_sensor', + CameraInfo: 'camera', + ClimateInfo: 'climate', + CoverInfo: 'cover', + FanInfo: 'fan', + LightInfo: 'light', + SensorInfo: 'sensor', + SwitchInfo: 'switch', + TextSensorInfo: 'sensor', +} + @attr.s class RuntimeEntryData: @@ -33,6 +50,8 @@ class RuntimeEntryData: 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) + loaded_platforms = attr.ib(type=Set[str], factory=set) + platform_load_lock = attr.ib(type=asyncio.Lock, factory=asyncio.Lock) def async_update_entity(self, hass: HomeAssistantType, component_key: str, key: int) -> None: @@ -48,9 +67,33 @@ class RuntimeEntryData: 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: + async def _ensure_platforms_loaded(self, hass: HomeAssistantType, + entry: ConfigEntry, + platforms: Set[str]): + async with self.platform_load_lock: + needed = platforms - self.loaded_platforms + tasks = [] + for platform in needed: + tasks.append(hass.config_entries.async_forward_entry_setup( + entry, platform)) + if tasks: + await asyncio.wait(tasks) + self.loaded_platforms |= needed + + async def async_update_static_infos( + self, hass: HomeAssistantType, entry: ConfigEntry, + infos: List[EntityInfo]) -> None: """Distribute an update of static infos to all platforms.""" + # First, load all platforms + needed_platforms = set() + for info in infos: + for info_type, platform in INFO_TYPE_TO_PLATFORM.items(): + if isinstance(info, info_type): + needed_platforms.add(platform) + break + await self._ensure_platforms_loaded(hass, entry, needed_platforms) + + # Then send dispatcher event signal = DISPATCHER_ON_LIST.format(entry_id=self.entry_id) async_dispatcher_send(hass, signal, infos)