diff --git a/homeassistant/components/google_mail/__init__.py b/homeassistant/components/google_mail/__init__.py index 441ecd3841f..7fae5f18da5 100644 --- a/homeassistant/components/google_mail/__init__.py +++ b/homeassistant/components/google_mail/__init__.py @@ -16,6 +16,8 @@ from .api import AsyncConfigEntryAuth from .const import DATA_AUTH, DATA_HASS_CONFIG, DOMAIN from .services import async_setup_services +type GoogleMailConfigEntry = ConfigEntry[AsyncConfigEntryAuth] + PLATFORMS = [Platform.NOTIFY, Platform.SENSOR] CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) @@ -28,13 +30,13 @@ 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: GoogleMailConfigEntry) -> bool: """Set up Google Mail from a config entry.""" implementation = await async_get_config_entry_implementation(hass, entry) session = OAuth2Session(hass, entry, implementation) auth = AsyncConfigEntryAuth(hass, session) await auth.check_and_refresh_token() - hass.data[DOMAIN][entry.entry_id] = auth + entry.runtime_data = auth hass.async_create_task( discovery.async_load_platform( @@ -55,10 +57,8 @@ 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: GoogleMailConfigEntry) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - hass.data[DOMAIN].pop(entry.entry_id) loaded_entries = [ entry for entry in hass.config_entries.async_entries(DOMAIN) @@ -68,4 +68,4 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: for service_name in hass.services.async_services_for_domain(DOMAIN): hass.services.async_remove(DOMAIN, service_name) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/google_mail/config_flow.py b/homeassistant/components/google_mail/config_flow.py index 5b5c760628b..5c81f7d49f5 100644 --- a/homeassistant/components/google_mail/config_flow.py +++ b/homeassistant/components/google_mail/config_flow.py @@ -9,10 +9,11 @@ from typing import Any, cast from google.oauth2.credentials import Credentials from googleapiclient.discovery import build -from homeassistant.config_entries import ConfigEntry, ConfigFlowResult +from homeassistant.config_entries import ConfigFlowResult from homeassistant.const import CONF_ACCESS_TOKEN, CONF_TOKEN from homeassistant.helpers import config_entry_oauth2_flow +from . import GoogleMailConfigEntry from .const import DEFAULT_ACCESS, DOMAIN @@ -23,7 +24,7 @@ class OAuth2FlowHandler( DOMAIN = DOMAIN - reauth_entry: ConfigEntry | None = None + reauth_entry: GoogleMailConfigEntry | None = None @property def logger(self) -> logging.Logger: diff --git a/homeassistant/components/google_mail/sensor.py b/homeassistant/components/google_mail/sensor.py index 1de72632de1..c832104d719 100644 --- a/homeassistant/components/google_mail/sensor.py +++ b/homeassistant/components/google_mail/sensor.py @@ -11,11 +11,10 @@ from homeassistant.components.sensor import ( SensorEntity, SensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import GoogleMailConfigEntry from .entity import GoogleMailEntity SCAN_INTERVAL = timedelta(minutes=15) @@ -28,12 +27,12 @@ SENSOR_TYPE = SensorEntityDescription( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: GoogleMailConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Google Mail sensor.""" - async_add_entities( - [GoogleMailSensor(hass.data[DOMAIN][entry.entry_id], SENSOR_TYPE)], True - ) + async_add_entities([GoogleMailSensor(entry.runtime_data, SENSOR_TYPE)], True) class GoogleMailSensor(GoogleMailEntity, SensorEntity): diff --git a/homeassistant/components/google_mail/services.py b/homeassistant/components/google_mail/services.py index e07e2be2101..2a81f7e6c51 100644 --- a/homeassistant/components/google_mail/services.py +++ b/homeassistant/components/google_mail/services.py @@ -3,16 +3,15 @@ from __future__ import annotations from datetime import datetime, timedelta +from typing import TYPE_CHECKING from googleapiclient.http import HttpRequest import voluptuous as vol -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, ServiceCall from homeassistant.helpers import config_validation as cv from homeassistant.helpers.service import async_extract_config_entry_ids -from .api import AsyncConfigEntryAuth from .const import ( ATTR_ENABLED, ATTR_END, @@ -26,6 +25,9 @@ from .const import ( DOMAIN, ) +if TYPE_CHECKING: + from . import GoogleMailConfigEntry + SERVICE_SET_VACATION = "set_vacation" SERVICE_VACATION_SCHEMA = vol.All( @@ -47,7 +49,9 @@ SERVICE_VACATION_SCHEMA = vol.All( async def async_setup_services(hass: HomeAssistant) -> None: """Set up services for Google Mail integration.""" - async def extract_gmail_config_entries(call: ServiceCall) -> list[ConfigEntry]: + async def extract_gmail_config_entries( + call: ServiceCall, + ) -> list[GoogleMailConfigEntry]: return [ entry for entry_id in await async_extract_config_entry_ids(hass, call) @@ -57,10 +61,11 @@ async def async_setup_services(hass: HomeAssistant) -> None: async def gmail_service(call: ServiceCall) -> None: """Call Google Mail service.""" - auth: AsyncConfigEntryAuth for entry in await extract_gmail_config_entries(call): - if not (auth := hass.data[DOMAIN].get(entry.entry_id)): - raise ValueError(f"Config entry not loaded: {entry.entry_id}") + try: + auth = entry.runtime_data + except AttributeError as ex: + raise ValueError(f"Config entry not loaded: {entry.entry_id}") from ex service = await auth.get_resource() _settings = {