"""Engine component helper.""" from __future__ import annotations from collections.abc import Awaitable, Callable import logging from types import ModuleType from typing import Generic, Protocol, TypeVar from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.setup import async_prepare_setup_platform from . import discovery from .typing import ConfigType, DiscoveryInfoType class Engine: """Base class for Home Assistant engines.""" async def async_internal_added_to_hass(self) -> None: """Run when engine about to be added to Home Assistant. Not to be extended by integrations. """ async def async_added_to_hass(self) -> None: """Run when engine about to be added to Home Assistant.""" async def async_internal_will_remove_from_hass(self) -> None: """Prepare to remove the engine from Home Assistant. Not to be extended by integrations. """ async def async_will_remove_from_hass(self) -> None: """Prepare to remove the engine from Home Assistant.""" _EngineT_co = TypeVar("_EngineT_co", bound=Engine, covariant=True) class EnginePlatformModule(Protocol[_EngineT_co]): """Protocol type for engine platform modules.""" async def async_setup_entry( self, hass: HomeAssistant, entry: ConfigEntry, ) -> _EngineT_co: """Set up an integration platform from a config entry.""" async def async_setup_platform( self, hass: HomeAssistant, ) -> _EngineT_co: """Set up an integration platform async.""" class EngineComponent(Generic[_EngineT_co]): """Track engines for a component.""" def __init__( self, logger: logging.Logger, domain: str, hass: HomeAssistant, config: ConfigType, ) -> None: """Initialize the engine component.""" self.logger = logger self.domain = domain self.hass = hass self.config = config self._engines: dict[str, _EngineT_co] = {} @callback def async_get_engine(self, config_entry_id: str) -> _EngineT_co | None: """Return a wrapped engine.""" return self._engines.get(config_entry_id) @callback def async_get_engines(self) -> list[_EngineT_co]: """Return a wrapped engine.""" return list(self._engines.values()) @callback def async_setup_discovery(self) -> None: """Initialize the engine component discovery.""" async def async_platform_discovered( platform: str, info: DiscoveryInfoType | None ) -> None: """Handle for discovered platform.""" await self.async_setup_domain(platform) discovery.async_listen_platform( self.hass, self.domain, async_platform_discovered ) async def async_setup_domain(self, domain: str) -> bool: """Set up an integration.""" async def setup(platform: EnginePlatformModule[_EngineT_co]) -> _EngineT_co: return await platform.async_setup_platform(self.hass) return await self._async_do_setup(domain, domain, setup) async def async_setup_entry(self, config_entry: ConfigEntry) -> bool: """Set up a config entry.""" async def setup(platform: EnginePlatformModule[_EngineT_co]) -> _EngineT_co: return await platform.async_setup_entry(self.hass, config_entry) return await self._async_do_setup( config_entry.entry_id, config_entry.domain, setup ) async def _async_do_setup( self, key: str, platform_domain: str, get_setup_coro: Callable[[ModuleType], Awaitable[_EngineT_co]], ) -> bool: """Set up an entry.""" platform = await async_prepare_setup_platform( self.hass, self.config, self.domain, platform_domain ) if platform is None: return False if key in self._engines: raise ValueError("Config entry has already been setup!") try: engine = await get_setup_coro(platform) await engine.async_internal_added_to_hass() await engine.async_added_to_hass() except Exception: # pylint: disable=broad-except self.logger.exception( "Error getting engine for %s (%s)", key, platform_domain ) return False self._engines[key] = engine return True async def async_unload_domain(self, domain: str) -> bool: """Unload a domain.""" return await self._async_do_unload(domain) async def async_unload_entry(self, config_entry: ConfigEntry) -> bool: """Unload a config entry.""" return await self._async_do_unload(config_entry.entry_id) async def _async_do_unload(self, key: str) -> bool: """Unload an engine.""" if (engine := self._engines.pop(key, None)) is None: raise ValueError("Config entry was never loaded!") try: await engine.async_internal_will_remove_from_hass() await engine.async_will_remove_from_hass() except Exception: # pylint: disable=broad-except self.logger.exception("Error unloading entry %s", key) return False return True