From bcf4c73f32d6b63e22376069c20c907e546cd900 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Fri, 19 Jul 2024 11:36:59 +0200 Subject: [PATCH] Migrate Wiz to config entry runtime data (#122091) --- homeassistant/components/wiz/__init__.py | 9 ++++----- homeassistant/components/wiz/binary_sensor.py | 11 +++++------ homeassistant/components/wiz/diagnostics.py | 9 +++------ homeassistant/components/wiz/light.py | 10 ++++------ homeassistant/components/wiz/number.py | 12 ++++++------ homeassistant/components/wiz/sensor.py | 13 ++++++------- homeassistant/components/wiz/switch.py | 10 ++++------ 7 files changed, 32 insertions(+), 42 deletions(-) diff --git a/homeassistant/components/wiz/__init__.py b/homeassistant/components/wiz/__init__.py index 79c317f178b..1bf3188e9e9 100644 --- a/homeassistant/components/wiz/__init__.py +++ b/homeassistant/components/wiz/__init__.py @@ -31,6 +31,8 @@ from .const import ( from .discovery import async_discover_devices, async_trigger_discovery from .models import WizData +type WizConfigEntry = ConfigEntry[WizData] + _LOGGER = logging.getLogger(__name__) PLATFORMS = [ @@ -135,9 +137,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await bulb.start_push(_async_push_update) bulb.set_discovery_callback(lambda bulb: async_trigger_discovery(hass, [bulb])) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = WizData( - coordinator=coordinator, bulb=bulb, scenes=scenes - ) + entry.runtime_data = WizData(coordinator=coordinator, bulb=bulb, scenes=scenes) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) entry.async_on_unload(entry.add_update_listener(_async_update_listener)) @@ -147,6 +147,5 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - data: WizData = hass.data[DOMAIN].pop(entry.entry_id) - await data.bulb.async_close() + await entry.runtime_data.bulb.async_close() return unload_ok diff --git a/homeassistant/components/wiz/binary_sensor.py b/homeassistant/components/wiz/binary_sensor.py index b58e120a9dd..3411ee200b9 100644 --- a/homeassistant/components/wiz/binary_sensor.py +++ b/homeassistant/components/wiz/binary_sensor.py @@ -10,13 +10,13 @@ from homeassistant.components.binary_sensor import ( BinarySensorDeviceClass, BinarySensorEntity, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import entity_registry as er from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback +from . import WizConfigEntry from .const import DOMAIN, SIGNAL_WIZ_PIR from .entity import WizEntity from .models import WizData @@ -26,17 +26,16 @@ OCCUPANCY_UNIQUE_ID = "{}_occupancy" async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WizConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the WiZ binary sensor platform.""" - wiz_data: WizData = hass.data[DOMAIN][entry.entry_id] - mac = wiz_data.bulb.mac + mac = entry.runtime_data.bulb.mac if er.async_get(hass).async_get_entity_id( Platform.BINARY_SENSOR, DOMAIN, OCCUPANCY_UNIQUE_ID.format(mac) ): - async_add_entities([WizOccupancyEntity(wiz_data, entry.title)]) + async_add_entities([WizOccupancyEntity(entry.runtime_data, entry.title)]) return cancel_dispatcher: Callable[[], None] | None = None @@ -47,7 +46,7 @@ async def async_setup_entry( assert cancel_dispatcher is not None cancel_dispatcher() cancel_dispatcher = None - async_add_entities([WizOccupancyEntity(wiz_data, entry.title)]) + async_add_entities([WizOccupancyEntity(entry.runtime_data, entry.title)]) cancel_dispatcher = async_dispatcher_connect( hass, SIGNAL_WIZ_PIR.format(mac), _async_add_occupancy_sensor diff --git a/homeassistant/components/wiz/diagnostics.py b/homeassistant/components/wiz/diagnostics.py index 5f617ebafe9..c58751c7fc0 100644 --- a/homeassistant/components/wiz/diagnostics.py +++ b/homeassistant/components/wiz/diagnostics.py @@ -5,24 +5,21 @@ from __future__ import annotations from typing import Any from homeassistant.components.diagnostics import async_redact_data -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .models import WizData +from . import WizConfigEntry TO_REDACT = {"roomId", "homeId"} async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: WizConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - wiz_data: WizData = hass.data[DOMAIN][entry.entry_id] return { "entry": { "title": entry.title, "data": dict(entry.data), }, - "data": async_redact_data(wiz_data.bulb.diagnostics, TO_REDACT), + "data": async_redact_data(entry.runtime_data.bulb.diagnostics, TO_REDACT), } diff --git a/homeassistant/components/wiz/light.py b/homeassistant/components/wiz/light.py index aece184720d..a3f36d580d2 100644 --- a/homeassistant/components/wiz/light.py +++ b/homeassistant/components/wiz/light.py @@ -19,7 +19,6 @@ from homeassistant.components.light import ( LightEntityFeature, filter_supported_color_modes, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.util.color import ( @@ -27,7 +26,7 @@ from homeassistant.util.color import ( color_temperature_mired_to_kelvin, ) -from .const import DOMAIN +from . import WizConfigEntry from .entity import WizToggleEntity from .models import WizData @@ -61,13 +60,12 @@ def _async_pilot_builder(**kwargs: Any) -> PilotBuilder: async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WizConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the WiZ Platform from config_flow.""" - wiz_data: WizData = hass.data[DOMAIN][entry.entry_id] - if wiz_data.bulb.bulbtype.bulb_type != BulbClass.SOCKET: - async_add_entities([WizBulbEntity(wiz_data, entry.title)]) + if entry.runtime_data.bulb.bulbtype.bulb_type != BulbClass.SOCKET: + async_add_entities([WizBulbEntity(entry.runtime_data, entry.title)]) class WizBulbEntity(WizToggleEntity, LightEntity): diff --git a/homeassistant/components/wiz/number.py b/homeassistant/components/wiz/number.py index 46708ac001e..0591e854d7d 100644 --- a/homeassistant/components/wiz/number.py +++ b/homeassistant/components/wiz/number.py @@ -13,12 +13,11 @@ from homeassistant.components.number import ( NumberEntityDescription, NumberMode, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import WizConfigEntry from .entity import WizEntity from .models import WizData @@ -68,15 +67,16 @@ NUMBERS: tuple[WizNumberEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WizConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the wiz speed number.""" - wiz_data: WizData = hass.data[DOMAIN][entry.entry_id] async_add_entities( - WizSpeedNumber(wiz_data, entry.title, description) + WizSpeedNumber(entry.runtime_data, entry.title, description) for description in NUMBERS - if getattr(wiz_data.bulb.bulbtype.features, description.required_feature) + if getattr( + entry.runtime_data.bulb.bulbtype.features, description.required_feature + ) ) diff --git a/homeassistant/components/wiz/sensor.py b/homeassistant/components/wiz/sensor.py index aae443e60d0..eb77686a5cf 100644 --- a/homeassistant/components/wiz/sensor.py +++ b/homeassistant/components/wiz/sensor.py @@ -8,7 +8,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( SIGNAL_STRENGTH_DECIBELS_MILLIWATT, EntityCategory, @@ -17,7 +16,7 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import WizConfigEntry from .entity import WizEntity from .models import WizData @@ -45,18 +44,18 @@ POWER_SENSORS: tuple[SensorEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WizConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the wiz sensor.""" - wiz_data: WizData = hass.data[DOMAIN][entry.entry_id] entities = [ - WizSensor(wiz_data, entry.title, description) for description in SENSORS + WizSensor(entry.runtime_data, entry.title, description) + for description in SENSORS ] - if wiz_data.coordinator.data is not None: + if entry.runtime_data.coordinator.data is not None: entities.extend( [ - WizPowerSensor(wiz_data, entry.title, description) + WizPowerSensor(entry.runtime_data, entry.title, description) for description in POWER_SENSORS ] ) diff --git a/homeassistant/components/wiz/switch.py b/homeassistant/components/wiz/switch.py index d94bf12da9f..4c089d2d6d2 100644 --- a/homeassistant/components/wiz/switch.py +++ b/homeassistant/components/wiz/switch.py @@ -8,24 +8,22 @@ from pywizlight import PilotBuilder from pywizlight.bulblibrary import BulbClass from homeassistant.components.switch import SwitchEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import WizConfigEntry from .entity import WizToggleEntity from .models import WizData async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WizConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the WiZ switch platform.""" - wiz_data: WizData = hass.data[DOMAIN][entry.entry_id] - if wiz_data.bulb.bulbtype.bulb_type == BulbClass.SOCKET: - async_add_entities([WizSocketEntity(wiz_data, entry.title)]) + if entry.runtime_data.bulb.bulbtype.bulb_type == BulbClass.SOCKET: + async_add_entities([WizSocketEntity(entry.runtime_data, entry.title)]) class WizSocketEntity(WizToggleEntity, SwitchEntity):