diff --git a/homeassistant/components/watttime/__init__.py b/homeassistant/components/watttime/__init__.py index 6e67994b11a2..91da8b477080 100644 --- a/homeassistant/components/watttime/__init__.py +++ b/homeassistant/components/watttime/__init__.py @@ -5,19 +5,18 @@ from __future__ import annotations from aiowatttime import Client from aiowatttime.errors import InvalidCredentialsError, WattTimeError -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed from homeassistant.helpers import aiohttp_client -from .const import DOMAIN, LOGGER -from .coordinator import WattTimeCoordinator +from .const import LOGGER +from .coordinator import WattTimeConfigEntry, WattTimeCoordinator PLATFORMS: list[Platform] = [Platform.SENSOR] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: WattTimeConfigEntry) -> bool: """Set up WattTime from a config entry.""" session = aiohttp_client.async_get_clientsession(hass) @@ -34,8 +33,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: coordinator = WattTimeCoordinator(hass, entry, client) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {}) - hass.data[DOMAIN][entry.entry_id] = coordinator + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) @@ -44,15 +42,13 @@ 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: WattTimeConfigEntry) -> bool: """Unload a config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) -async def async_reload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> None: +async def async_reload_entry( + hass: HomeAssistant, config_entry: WattTimeConfigEntry +) -> None: """Handle an options update.""" await hass.config_entries.async_reload(config_entry.entry_id) diff --git a/homeassistant/components/watttime/config_flow.py b/homeassistant/components/watttime/config_flow.py index ad676e166c58..7587611457f0 100644 --- a/homeassistant/components/watttime/config_flow.py +++ b/homeassistant/components/watttime/config_flow.py @@ -9,12 +9,7 @@ from aiowatttime import Client from aiowatttime.errors import CoordinatesNotFoundError, InvalidCredentialsError import voluptuous as vol -from homeassistant.config_entries import ( - ConfigEntry, - ConfigFlow, - ConfigFlowResult, - OptionsFlow, -) +from homeassistant.config_entries import ConfigFlow, ConfigFlowResult, OptionsFlow from homeassistant.const import ( CONF_LATITUDE, CONF_LONGITUDE, @@ -31,6 +26,7 @@ from .const import ( DOMAIN, LOGGER, ) +from .coordinator import WattTimeConfigEntry CONF_LOCATION_TYPE = "location_type" @@ -127,7 +123,7 @@ class WattTimeConfigFlow(ConfigFlow, domain=DOMAIN): @staticmethod @callback def async_get_options_flow( - config_entry: ConfigEntry, + config_entry: WattTimeConfigEntry, ) -> WattTimeOptionsFlowHandler: """Define the config flow to handle options.""" return WattTimeOptionsFlowHandler() diff --git a/homeassistant/components/watttime/coordinator.py b/homeassistant/components/watttime/coordinator.py index a726555db538..6caf7e415a68 100644 --- a/homeassistant/components/watttime/coordinator.py +++ b/homeassistant/components/watttime/coordinator.py @@ -18,16 +18,18 @@ from .const import DOMAIN, LOGGER DEFAULT_UPDATE_INTERVAL = timedelta(minutes=5) +type WattTimeConfigEntry = ConfigEntry[WattTimeCoordinator] + class WattTimeCoordinator(DataUpdateCoordinator[RealTimeEmissionsResponseType]): """Coordinator for WattTime data updates.""" - config_entry: ConfigEntry + config_entry: WattTimeConfigEntry def __init__( self, hass: HomeAssistant, - entry: ConfigEntry, + entry: WattTimeConfigEntry, client: Client, ) -> None: """Initialize the coordinator.""" diff --git a/homeassistant/components/watttime/diagnostics.py b/homeassistant/components/watttime/diagnostics.py index b779b2759d1d..f215cc1a680e 100644 --- a/homeassistant/components/watttime/diagnostics.py +++ b/homeassistant/components/watttime/diagnostics.py @@ -5,7 +5,6 @@ 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_LATITUDE, CONF_LONGITUDE, @@ -15,8 +14,8 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant -from .const import CONF_BALANCING_AUTHORITY, CONF_BALANCING_AUTHORITY_ABBREV, DOMAIN -from .coordinator import WattTimeCoordinator +from .const import CONF_BALANCING_AUTHORITY, CONF_BALANCING_AUTHORITY_ABBREV +from .coordinator import WattTimeConfigEntry CONF_TITLE = "title" @@ -34,15 +33,13 @@ TO_REDACT = { async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: WattTimeConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: WattTimeCoordinator = hass.data[DOMAIN][entry.entry_id] - return async_redact_data( { "entry": entry.as_dict(), - "data": coordinator.data, + "data": entry.runtime_data.data, }, TO_REDACT, ) diff --git a/homeassistant/components/watttime/sensor.py b/homeassistant/components/watttime/sensor.py index 23824a1369a0..3dfa28132bb1 100644 --- a/homeassistant/components/watttime/sensor.py +++ b/homeassistant/components/watttime/sensor.py @@ -10,7 +10,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( ATTR_LATITUDE, ATTR_LONGITUDE, @@ -25,7 +24,7 @@ from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import CONF_BALANCING_AUTHORITY, CONF_BALANCING_AUTHORITY_ABBREV, DOMAIN -from .coordinator import WattTimeCoordinator +from .coordinator import WattTimeConfigEntry, WattTimeCoordinator ATTR_BALANCING_AUTHORITY = "balancing_authority" @@ -51,11 +50,11 @@ REALTIME_EMISSIONS_SENSOR_DESCRIPTIONS = ( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WattTimeConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up WattTime sensors based on a config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( [ RealtimeEmissionsSensor(coordinator, entry, description) @@ -73,7 +72,7 @@ class RealtimeEmissionsSensor(CoordinatorEntity[WattTimeCoordinator], SensorEnti def __init__( self, coordinator: WattTimeCoordinator, - entry: ConfigEntry, + entry: WattTimeConfigEntry, description: SensorEntityDescription, ) -> None: """Initialize the sensor."""