diff --git a/homeassistant/components/baf/__init__.py b/homeassistant/components/baf/__init__.py index d3b29b52e44..659cb10eba1 100644 --- a/homeassistant/components/baf/__init__.py +++ b/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) diff --git a/homeassistant/components/baf/binary_sensor.py b/homeassistant/components/baf/binary_sensor.py index e95e197b8be..b1076a99f8a 100644 --- a/homeassistant/components/baf/binary_sensor.py +++ b/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) diff --git a/homeassistant/components/baf/climate.py b/homeassistant/components/baf/climate.py index f451c5e7a71..38407813d37 100644 --- a/homeassistant/components/baf/climate.py +++ b/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): diff --git a/homeassistant/components/baf/fan.py b/homeassistant/components/baf/fan.py index 6c90e2a53cb..d8c800ea512 100644 --- a/homeassistant/components/baf/fan.py +++ b/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): diff --git a/homeassistant/components/baf/light.py b/homeassistant/components/baf/light.py index e203e12cf96..2fb36ed874f 100644 --- a/homeassistant/components/baf/light.py +++ b/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): diff --git a/homeassistant/components/baf/models.py b/homeassistant/components/baf/models.py index c94b73d9abd..3bb574d5a19 100644 --- a/homeassistant/components/baf/models.py +++ b/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: diff --git a/homeassistant/components/baf/number.py b/homeassistant/components/baf/number.py index 43da381391c..bf9e837eea1 100644 --- a/homeassistant/components/baf/number.py +++ b/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) diff --git a/homeassistant/components/baf/sensor.py b/homeassistant/components/baf/sensor.py index fc052b1e48b..a97e2945564 100644 --- a/homeassistant/components/baf/sensor.py +++ b/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 diff --git a/homeassistant/components/baf/switch.py b/homeassistant/components/baf/switch.py index 38248e48d09..789ea365d6d 100644 --- a/homeassistant/components/baf/switch.py +++ b/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: