diff --git a/homeassistant/components/justnimbus/__init__.py b/homeassistant/components/justnimbus/__init__.py index 123807d887c..5f369027b00 100644 --- a/homeassistant/components/justnimbus/__init__.py +++ b/homeassistant/components/justnimbus/__init__.py @@ -2,15 +2,14 @@ from __future__ import annotations -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed -from .const import DOMAIN, PLATFORMS -from .coordinator import JustNimbusCoordinator +from .const import PLATFORMS +from .coordinator import JustNimbusConfigEntry, JustNimbusCoordinator -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: JustNimbusConfigEntry) -> bool: """Set up JustNimbus from a config entry.""" if "zip_code" in entry.data: coordinator = JustNimbusCoordinator(hass, entry) @@ -18,13 +17,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: raise ConfigEntryAuthFailed await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(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: JustNimbusConfigEntry) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - hass.data[DOMAIN].pop(entry.entry_id) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/justnimbus/coordinator.py b/homeassistant/components/justnimbus/coordinator.py index a6945c45417..b51058a8e54 100644 --- a/homeassistant/components/justnimbus/coordinator.py +++ b/homeassistant/components/justnimbus/coordinator.py @@ -16,13 +16,17 @@ from .const import CONF_ZIP_CODE, DOMAIN _LOGGER = logging.getLogger(__name__) +type JustNimbusConfigEntry = ConfigEntry[JustNimbusCoordinator] + class JustNimbusCoordinator(DataUpdateCoordinator[justnimbus.JustNimbusModel]): """Data update coordinator.""" - config_entry: ConfigEntry + config_entry: JustNimbusConfigEntry - def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None: + def __init__( + self, hass: HomeAssistant, config_entry: JustNimbusConfigEntry + ) -> None: """Initialize the coordinator.""" super().__init__( hass, diff --git a/homeassistant/components/justnimbus/sensor.py b/homeassistant/components/justnimbus/sensor.py index 1e288e272cd..88f12cad113 100644 --- a/homeassistant/components/justnimbus/sensor.py +++ b/homeassistant/components/justnimbus/sensor.py @@ -12,7 +12,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONF_CLIENT_ID, EntityCategory, @@ -24,8 +23,7 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.typing import StateType -from . import JustNimbusCoordinator -from .const import DOMAIN +from .coordinator import JustNimbusConfigEntry, JustNimbusCoordinator from .entity import JustNimbusEntity @@ -102,16 +100,15 @@ SENSOR_TYPES = ( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: JustNimbusConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up the JustNimbus sensor.""" - coordinator: JustNimbusCoordinator = hass.data[DOMAIN][entry.entry_id] async_add_entities( JustNimbusSensor( device_id=entry.data[CONF_CLIENT_ID], description=description, - coordinator=coordinator, + coordinator=entry.runtime_data, ) for description in SENSOR_TYPES )