From bb2883a5a831c824bc9ff0d1369516f2c8d99a59 Mon Sep 17 00:00:00 2001 From: Rami Mosleh Date: Thu, 13 Jun 2024 17:58:05 +0300 Subject: [PATCH] Store imap coordinator in runtime_data (#119611) --- homeassistant/components/imap/__init__.py | 27 +++++++++----------- homeassistant/components/imap/diagnostics.py | 10 +++----- homeassistant/components/imap/sensor.py | 17 +++++------- 3 files changed, 22 insertions(+), 32 deletions(-) diff --git a/homeassistant/components/imap/__init__.py b/homeassistant/components/imap/__init__.py index f39a78925c1..f62edf1451f 100644 --- a/homeassistant/components/imap/__init__.py +++ b/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 diff --git a/homeassistant/components/imap/diagnostics.py b/homeassistant/components/imap/diagnostics.py index 8afe3e327ba..d402053520a 100644 --- a/homeassistant/components/imap/diagnostics.py +++ b/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, diff --git a/homeassistant/components/imap/sensor.py b/homeassistant/components/imap/sensor.py index 0a9070d7a5e..625af9ce6a1 100644 --- a/homeassistant/components/imap/sensor.py +++ b/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."""