ESPHome load platforms lazily (#24594)

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))
+            entry, platform))
-    await asyncio.wait(tasks)
+    if tasks:
-
+        await asyncio.wait(tasks)
     return True
 
 

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,
+    async def _ensure_platforms_loaded(self, hass: HomeAssistantType,
-                                  infos: List[EntityInfo]) -> None:
+                                       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)