Use runtime_data in watttime integration (#168630)

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)

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()

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."""

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,
     )

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."""