diff --git a/homeassistant/components/weatherkit/__init__.py b/homeassistant/components/weatherkit/__init__.py index 4cbac2b32d83..01f53ffdf430 100644 --- a/homeassistant/components/weatherkit/__init__.py +++ b/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) diff --git a/homeassistant/components/weatherkit/coordinator.py b/homeassistant/components/weatherkit/coordinator.py index fd790ee230f5..2575f174f86f 100644 --- a/homeassistant/components/weatherkit/coordinator.py +++ b/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.""" diff --git a/homeassistant/components/weatherkit/sensor.py b/homeassistant/components/weatherkit/sensor.py index b3639fa53561..224f59864779 100644 --- a/homeassistant/components/weatherkit/sensor.py +++ b/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 diff --git a/homeassistant/components/weatherkit/weather.py b/homeassistant/components/weatherkit/weather.py index b57e488d06a1..0534c4207854 100644 --- a/homeassistant/components/weatherkit/weather.py +++ b/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 = {