diff --git a/homeassistant/components/weatherflow/__init__.py b/homeassistant/components/weatherflow/__init__.py index 3e30d15aebee..9851062c5cd1 100644 --- a/homeassistant/components/weatherflow/__init__.py +++ b/homeassistant/components/weatherflow/__init__.py @@ -21,8 +21,10 @@ PLATFORMS = [ Platform.SENSOR, ] +type WeatherFlowConfigEntry = ConfigEntry[WeatherFlowListener] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: WeatherFlowConfigEntry) -> bool: """Set up WeatherFlow from a config entry.""" client = WeatherFlowListener() @@ -56,7 +58,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: except ListenerError as ex: raise ConfigEntryNotReady from ex - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = client + entry.runtime_data = client await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) async def _async_handle_ha_shutdown(event: Event) -> None: @@ -70,21 +72,23 @@ 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: WeatherFlowConfigEntry +) -> bool: """Unload a config entry.""" if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - client: WeatherFlowListener = hass.data[DOMAIN].pop(entry.entry_id, None) - if client: - await client.stop_listening() + await entry.runtime_data.stop_listening() return unload_ok async def async_remove_config_entry_device( - hass: HomeAssistant, config_entry: ConfigEntry, device_entry: DeviceEntry + hass: HomeAssistant, + config_entry: WeatherFlowConfigEntry, + device_entry: DeviceEntry, ) -> bool: """Remove a config entry from a device.""" - client: WeatherFlowListener = hass.data[DOMAIN][config_entry.entry_id] + client = config_entry.runtime_data return not any( identifier for identifier in device_entry.identifiers diff --git a/homeassistant/components/weatherflow/event.py b/homeassistant/components/weatherflow/event.py index 05f7ecc28651..7e2d3122249f 100644 --- a/homeassistant/components/weatherflow/event.py +++ b/homeassistant/components/weatherflow/event.py @@ -7,12 +7,12 @@ from dataclasses import dataclass from pyweatherflowudp.device import EVENT_RAIN_START, EVENT_STRIKE, WeatherFlowDevice from homeassistant.components.event import EventEntity, EventEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback +from . import WeatherFlowConfigEntry from .const import DOMAIN, LOGGER, format_dispatch_call @@ -42,7 +42,7 @@ EVENT_DESCRIPTIONS: list[WeatherFlowEventEntityDescription] = [ async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: WeatherFlowConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up WeatherFlow event entities using config entry.""" diff --git a/homeassistant/components/weatherflow/sensor.py b/homeassistant/components/weatherflow/sensor.py index 3d4881324ba8..d9a1486b415d 100644 --- a/homeassistant/components/weatherflow/sensor.py +++ b/homeassistant/components/weatherflow/sensor.py @@ -22,7 +22,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( DEGREE, LIGHT_LUX, @@ -46,6 +45,7 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.typing import StateType from homeassistant.util.unit_system import METRIC_SYSTEM +from . import WeatherFlowConfigEntry from .const import DOMAIN, LOGGER, format_dispatch_call @@ -295,7 +295,7 @@ SENSORS: tuple[WeatherFlowSensorEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: WeatherFlowConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up WeatherFlow sensors using config entry."""