Use runtime_data in WeatherKit (#168682)

homeassistant/components/weatherkit/__init__.py

@@ -8,28 +8,19 @@ from apple_weatherkit.client import (
     WeatherKitApiClientError,
 )
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import (
-    CONF_KEY_ID,
-    CONF_KEY_PEM,
-    CONF_SERVICE_ID,
-    CONF_TEAM_ID,
-    DOMAIN,
-    LOGGER,
-)
-from .coordinator import WeatherKitDataUpdateCoordinator
+from .const import CONF_KEY_ID, CONF_KEY_PEM, CONF_SERVICE_ID, CONF_TEAM_ID, LOGGER
+from .coordinator import WeatherKitConfigEntry, WeatherKitDataUpdateCoordinator
 
 PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.WEATHER]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: WeatherKitConfigEntry) -> bool:
     """Set up this integration using UI."""
-    hass.data.setdefault(DOMAIN, {})
     coordinator = WeatherKitDataUpdateCoordinator(
         hass=hass,
         config_entry=entry,
@@ -51,14 +42,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady from ex
 
     await coordinator.async_config_entry_first_refresh()
-    hass.data[DOMAIN][entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: WeatherKitConfigEntry) -> bool:
     """Handle removal of an entry."""
-    if unloaded := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unloaded
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/weatherkit/coordinator.py

@@ -25,18 +25,20 @@ STALE_DATA_THRESHOLD = timedelta(hours=1)
 
 HOURLY_FORECAST_DURATION = timedelta(days=7)
 
+type WeatherKitConfigEntry = ConfigEntry[WeatherKitDataUpdateCoordinator]
+
 
 class WeatherKitDataUpdateCoordinator(DataUpdateCoordinator):
     """Class to manage fetching data from the API."""
 
-    config_entry: ConfigEntry
+    config_entry: WeatherKitConfigEntry
     supported_data_sets: list[DataSetType] | None = None
     last_updated_at: datetime | None = None
 
     def __init__(
         self,
         hass: HomeAssistant,
-        config_entry: ConfigEntry,
+        config_entry: WeatherKitConfigEntry,
         client: WeatherKitApiClient,
     ) -> None:
         """Initialize."""

homeassistant/components/weatherkit/sensor.py

@@ -6,15 +6,14 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfVolumetricFlux
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import ATTR_CURRENT_WEATHER, DOMAIN
-from .coordinator import WeatherKitDataUpdateCoordinator
+from .const import ATTR_CURRENT_WEATHER
+from .coordinator import WeatherKitConfigEntry, WeatherKitDataUpdateCoordinator
 from .entity import WeatherKitEntity
 
 SENSORS = (
@@ -35,13 +34,11 @@ SENSORS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: WeatherKitConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Add sensor entities from a config_entry."""
-    coordinator: WeatherKitDataUpdateCoordinator = hass.data[DOMAIN][
-        config_entry.entry_id
-    ]
+    coordinator = config_entry.runtime_data
 
     async_add_entities(
         WeatherKitSensor(coordinator, description) for description in SENSORS

homeassistant/components/weatherkit/weather.py

@@ -21,7 +21,6 @@ from homeassistant.components.weather import (
     SingleCoordinatorWeatherEntity,
     WeatherEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     UnitOfLength,
     UnitOfPressure,
@@ -36,23 +35,18 @@ from .const import (
     ATTR_FORECAST_DAILY,
     ATTR_FORECAST_HOURLY,
     ATTRIBUTION,
-    DOMAIN,
 )
-from .coordinator import WeatherKitDataUpdateCoordinator
+from .coordinator import WeatherKitConfigEntry, WeatherKitDataUpdateCoordinator
 from .entity import WeatherKitEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: WeatherKitConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Add a weather entity from a config_entry."""
-    coordinator: WeatherKitDataUpdateCoordinator = hass.data[DOMAIN][
-        config_entry.entry_id
-    ]
-
-    async_add_entities([WeatherKitWeather(coordinator)])
+    async_add_entities([WeatherKitWeather(config_entry.runtime_data)])
 
 
 condition_code_to_hass = {