Store imap coordinator in runtime_data (#119611)

homeassistant/components/imap/__init__.py

@@ -4,12 +4,11 @@ from __future__ import annotations
 
 import asyncio
 import logging
-from typing import TYPE_CHECKING
 
 from aioimaplib import IMAP4_SSL, AioImapException, Response
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
 from homeassistant.core import (
     HomeAssistant,
@@ -29,6 +28,7 @@ from homeassistant.helpers.typing import ConfigType
 
 from .const import CONF_ENABLE_PUSH, DOMAIN
 from .coordinator import (
+    ImapDataUpdateCoordinator,
     ImapMessage,
     ImapPollingDataUpdateCoordinator,
     ImapPushDataUpdateCoordinator,
@@ -65,17 +65,18 @@ SERVICE_MOVE_SCHEMA = _SERVICE_UID_SCHEMA.extend(
 SERVICE_DELETE_SCHEMA = _SERVICE_UID_SCHEMA
 SERVICE_FETCH_TEXT_SCHEMA = _SERVICE_UID_SCHEMA
 
+type ImapConfigEntry = ConfigEntry[ImapDataUpdateCoordinator]
+
 
 async def async_get_imap_client(hass: HomeAssistant, entry_id: str) -> IMAP4_SSL:
     """Get IMAP client and connect."""
-    if hass.data[DOMAIN].get(entry_id) is None:
+    if (entry := hass.config_entries.async_get_entry(entry_id)) is None or (
+        entry.state is not ConfigEntryState.LOADED
+    ):
         raise ServiceValidationError(
             translation_domain=DOMAIN,
             translation_key="invalid_entry",
         )
-    entry = hass.config_entries.async_get_entry(entry_id)
-    if TYPE_CHECKING:
-        assert entry is not None
     try:
         client = await connect_to_server(entry.data)
     except InvalidAuth as exc:
@@ -235,7 +236,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: ImapConfigEntry) -> bool:
     """Set up imap from a config entry."""
     try:
         imap_client: IMAP4_SSL = await connect_to_server(dict(entry.data))
@@ -255,12 +256,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     else:
         coordinator_class = ImapPollingDataUpdateCoordinator
 
-    coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = (
-        coordinator_class(hass, imap_client, entry)
-    )
+    coordinator: ImapDataUpdateCoordinator = coordinator_class(hass, imap_client, entry)
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, coordinator.shutdown)
@@ -271,11 +270,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: ImapConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        coordinator: (
-            ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator
-        ) = hass.data[DOMAIN].pop(entry.entry_id)
+        coordinator = entry.runtime_data
         await coordinator.shutdown()
     return unload_ok

homeassistant/components/imap/diagnostics.py

@@ -5,18 +5,16 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant, callback
 
-from .const import DOMAIN
-from .coordinator import ImapDataUpdateCoordinator
+from . import ImapConfigEntry
 
 REDACT_CONFIG = {CONF_PASSWORD, CONF_USERNAME}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: ImapConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
     return _async_get_diagnostics(hass, entry)
@@ -25,11 +23,11 @@ async def async_get_config_entry_diagnostics(
 @callback
 def _async_get_diagnostics(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: ImapConfigEntry,
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
     redacted_config = async_redact_data(entry.data, REDACT_CONFIG)
-    coordinator: ImapDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     return {
         "config": redacted_config,

homeassistant/components/imap/sensor.py

@@ -7,15 +7,15 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import ImapPollingDataUpdateCoordinator, ImapPushDataUpdateCoordinator
+from . import ImapConfigEntry
 from .const import DOMAIN
+from .coordinator import ImapDataUpdateCoordinator
 
 IMAP_MAIL_COUNT_DESCRIPTION = SensorEntityDescription(
     key="imap_mail_count",
@@ -27,27 +27,22 @@ IMAP_MAIL_COUNT_DESCRIPTION = SensorEntityDescription(
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: ImapConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the Imap sensor."""
 
-    coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = (
-        hass.data[DOMAIN][entry.entry_id]
-    )
+    coordinator = entry.runtime_data
     async_add_entities([ImapSensor(coordinator, IMAP_MAIL_COUNT_DESCRIPTION)])
 
 
-class ImapSensor(
-    CoordinatorEntity[ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator],
-    SensorEntity,
-):
+class ImapSensor(CoordinatorEntity[ImapDataUpdateCoordinator], SensorEntity):
     """Representation of an IMAP sensor."""
 
     _attr_has_entity_name = True
 
     def __init__(
         self,
-        coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator,
+        coordinator: ImapDataUpdateCoordinator,
         description: SensorEntityDescription,
     ) -> None:
         """Initialize the sensor."""