Migrate baf to use config entry runtime_data (#117081)

homeassistant/components/baf/__init__.py

@@ -10,11 +10,13 @@ from aiobafi6.exceptions import DeviceUUIDMismatchError
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_IP_ADDRESS, Platform
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
 
-from .const import DOMAIN, QUERY_INTERVAL, RUN_TIMEOUT
-from .models import BAFData
+from .const import QUERY_INTERVAL, RUN_TIMEOUT
+
+BAFConfigEntry = ConfigEntry[Device]
+
 
 PLATFORMS: list[Platform] = [
     Platform.BINARY_SENSOR,
@@ -27,7 +29,7 @@ PLATFORMS: list[Platform] = [
 ]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: BAFConfigEntry) -> bool:
     """Set up Big Ass Fans from a config entry."""
     ip_address = entry.data[CONF_IP_ADDRESS]
 
@@ -46,16 +48,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         run_future.cancel()
         raise ConfigEntryNotReady(f"Timed out connecting to {ip_address}") from ex
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = BAFData(device, run_future)
+    @callback
+    def _async_cancel_run() -> None:
+        run_future.cancel()
+
+    entry.runtime_data = device
+    entry.async_on_unload(_async_cancel_run)
     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: BAFConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        data: BAFData = hass.data[DOMAIN].pop(entry.entry_id)
-        data.run_future.cancel()
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/baf/binary_sensor.py

@@ -8,7 +8,6 @@ from typing import cast
 
 from aiobafi6 import Device
 
-from homeassistant import config_entries
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
@@ -17,9 +16,8 @@ from homeassistant.components.binary_sensor import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import BAFConfigEntry
 from .entity import BAFEntity
-from .models import BAFData
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -42,12 +40,11 @@ OCCUPANCY_SENSORS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: config_entries.ConfigEntry,
+    entry: BAFConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up BAF binary sensors."""
-    data: BAFData = hass.data[DOMAIN][entry.entry_id]
-    device = data.device
+    device = entry.runtime_data
     sensors_descriptions: list[BAFBinarySensorDescription] = []
     if device.has_occupancy:
         sensors_descriptions.extend(OCCUPANCY_SENSORS)

homeassistant/components/baf/climate.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 from typing import Any
 
-from homeassistant import config_entries
 from homeassistant.components.climate import (
     ClimateEntity,
     ClimateEntityFeature,
@@ -15,20 +14,19 @@ from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import BAFConfigEntry
 from .entity import BAFEntity
-from .models import BAFData
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: config_entries.ConfigEntry,
+    entry: BAFConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up BAF fan auto comfort."""
-    data: BAFData = hass.data[DOMAIN][entry.entry_id]
-    if data.device.has_fan and data.device.has_auto_comfort:
-        async_add_entities([BAFAutoComfort(data.device)])
+    device = entry.runtime_data
+    if device.has_fan and device.has_auto_comfort:
+        async_add_entities([BAFAutoComfort(device)])
 
 
 class BAFAutoComfort(BAFEntity, ClimateEntity):

homeassistant/components/baf/fan.py

@@ -7,7 +7,6 @@ from typing import Any
 
 from aiobafi6 import OffOnAuto
 
-from homeassistant import config_entries
 from homeassistant.components.fan import (
     DIRECTION_FORWARD,
     DIRECTION_REVERSE,
@@ -21,20 +20,20 @@ from homeassistant.util.percentage import (
     ranged_value_to_percentage,
 )
 
-from .const import DOMAIN, PRESET_MODE_AUTO, SPEED_COUNT, SPEED_RANGE
+from . import BAFConfigEntry
+from .const import PRESET_MODE_AUTO, SPEED_COUNT, SPEED_RANGE
 from .entity import BAFEntity
-from .models import BAFData
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: config_entries.ConfigEntry,
+    entry: BAFConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up SenseME fans."""
-    data: BAFData = hass.data[DOMAIN][entry.entry_id]
-    if data.device.has_fan:
-        async_add_entities([BAFFan(data.device)])
+    device = entry.runtime_data
+    if device.has_fan:
+        async_add_entities([BAFFan(device)])
 
 
 class BAFFan(BAFEntity, FanEntity):

homeassistant/components/baf/light.py

@@ -6,7 +6,6 @@ from typing import Any
 
 from aiobafi6 import Device, OffOnAuto
 
-from homeassistant import config_entries
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
     ATTR_COLOR_TEMP,
@@ -20,21 +19,20 @@ from homeassistant.util.color import (
     color_temperature_mired_to_kelvin,
 )
 
-from .const import DOMAIN
+from . import BAFConfigEntry
 from .entity import BAFEntity
-from .models import BAFData
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: config_entries.ConfigEntry,
+    entry: BAFConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up BAF lights."""
-    data: BAFData = hass.data[DOMAIN][entry.entry_id]
-    if data.device.has_light:
-        klass = BAFFanLight if data.device.has_fan else BAFStandaloneLight
-        async_add_entities([klass(data.device)])
+    device = entry.runtime_data
+    if device.has_light:
+        klass = BAFFanLight if device.has_fan else BAFStandaloneLight
+        async_add_entities([klass(device)])
 
 
 class BAFLight(BAFEntity, LightEntity):

homeassistant/components/baf/models.py

@@ -2,19 +2,8 @@
 
 from __future__ import annotations
 
-import asyncio
 from dataclasses import dataclass
 
-from aiobafi6 import Device
-
-
-@dataclass
-class BAFData:
-    """Data for the baf integration."""
-
-    device: Device
-    run_future: asyncio.Future
-
 
 @dataclass
 class BAFDiscovery:

homeassistant/components/baf/number.py

@@ -8,7 +8,6 @@ from typing import cast
 
 from aiobafi6 import Device
 
-from homeassistant import config_entries
 from homeassistant.components.number import (
     NumberEntity,
     NumberEntityDescription,
@@ -18,9 +17,9 @@ from homeassistant.const import EntityCategory, UnitOfTime
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, HALF_DAY_SECS, ONE_DAY_SECS, ONE_MIN_SECS, SPEED_RANGE
+from . import BAFConfigEntry
+from .const import HALF_DAY_SECS, ONE_DAY_SECS, ONE_MIN_SECS, SPEED_RANGE
 from .entity import BAFEntity
-from .models import BAFData
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -116,12 +115,11 @@ LIGHT_NUMBER_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: config_entries.ConfigEntry,
+    entry: BAFConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up BAF numbers."""
-    data: BAFData = hass.data[DOMAIN][entry.entry_id]
-    device = data.device
+    device = entry.runtime_data
     descriptions: list[BAFNumberDescription] = []
     if device.has_fan:
         descriptions.extend(FAN_NUMBER_DESCRIPTIONS)

homeassistant/components/baf/sensor.py

@@ -8,7 +8,6 @@ from typing import cast
 
 from aiobafi6 import Device
 
-from homeassistant import config_entries
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
@@ -24,9 +23,8 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import BAFConfigEntry
 from .entity import BAFEntity
-from .models import BAFData
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -94,12 +92,11 @@ FAN_SENSORS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: config_entries.ConfigEntry,
+    entry: BAFConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up BAF fan sensors."""
-    data: BAFData = hass.data[DOMAIN][entry.entry_id]
-    device = data.device
+    device = entry.runtime_data
     sensors_descriptions: list[BAFSensorDescription] = [
         description
         for description in DEFINED_ONLY_SENSORS

homeassistant/components/baf/switch.py

@@ -8,15 +8,13 @@ from typing import Any, cast
 
 from aiobafi6 import Device
 
-from homeassistant import config_entries
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import BAFConfigEntry
 from .entity import BAFEntity
-from .models import BAFData
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -104,12 +102,11 @@ LIGHT_SWITCHES = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: config_entries.ConfigEntry,
+    entry: BAFConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up BAF fan switches."""
-    data: BAFData = hass.data[DOMAIN][entry.entry_id]
-    device = data.device
+    device = entry.runtime_data
     descriptions: list[BAFSwitchDescription] = []
     descriptions.extend(BASE_SWITCHES)
     if device.has_fan: