From f6ac856b8dcbed0b80f0e5051a0a78730cb749a6 Mon Sep 17 00:00:00 2001 From: "David F. Mulcahey" Date: Sat, 11 Dec 2021 11:50:03 -0500 Subject: [PATCH] Use async_on_unload from config entry in ZHA (#61015) * remove DATA_ZHA_DISPATCHERS * update typing information * fix rebase --- homeassistant/components/zha/__init__.py | 17 ++++++++-------- .../components/zha/alarm_control_panel.py | 13 ++++++++---- homeassistant/components/zha/binary_sensor.py | 13 ++++++++---- homeassistant/components/zha/climate.py | 13 ++++++++---- homeassistant/components/zha/core/const.py | 1 - homeassistant/components/zha/cover.py | 13 ++++++++---- .../components/zha/device_tracker.py | 13 ++++++++---- homeassistant/components/zha/fan.py | 20 +++++++++---------- homeassistant/components/zha/light.py | 13 ++++++++---- homeassistant/components/zha/lock.py | 12 +++++++---- homeassistant/components/zha/number.py | 13 ++++++++---- homeassistant/components/zha/sensor.py | 3 +-- homeassistant/components/zha/switch.py | 13 ++++++++---- 13 files changed, 99 insertions(+), 58 deletions(-) diff --git a/homeassistant/components/zha/__init__.py b/homeassistant/components/zha/__init__.py index d6578be775f..42f14e63927 100644 --- a/homeassistant/components/zha/__init__.py +++ b/homeassistant/components/zha/__init__.py @@ -12,6 +12,7 @@ from homeassistant.core import HomeAssistant import homeassistant.helpers.config_validation as cv from homeassistant.helpers.device_registry import CONNECTION_ZIGBEE from homeassistant.helpers.dispatcher import async_dispatcher_send +from homeassistant.helpers.typing import ConfigType from . import api from .core import ZHAGateway @@ -27,7 +28,6 @@ from .core.const import ( CONF_ZIGPY, DATA_ZHA, DATA_ZHA_CONFIG, - DATA_ZHA_DISPATCHERS, DATA_ZHA_GATEWAY, DATA_ZHA_PLATFORM_LOADED, DATA_ZHA_SHUTDOWN_TASK, @@ -72,7 +72,7 @@ CENTICELSIUS = "C-100" _LOGGER = logging.getLogger(__name__) -async def async_setup(hass, config): +async def async_setup(hass: HomeAssistant, config: ConfigType): """Set up ZHA from config.""" hass.data[DATA_ZHA] = {} @@ -83,7 +83,9 @@ async def async_setup(hass, config): return True -async def async_setup_entry(hass, config_entry): +async def async_setup_entry( + hass: HomeAssistant, config_entry: config_entries.ConfigEntry +): """Set up ZHA. Will automatically load components to support devices found on the network. @@ -101,7 +103,6 @@ async def async_setup_entry(hass, config_entry): zha_gateway = ZHAGateway(hass, config, config_entry) await zha_gateway.async_initialize() - zha_data[DATA_ZHA_DISPATCHERS] = [] zha_data[DATA_ZHA_PLATFORM_LOADED] = [] for platform in PLATFORMS: coro = hass.config_entries.async_forward_entry_setup(config_entry, platform) @@ -131,7 +132,9 @@ async def async_setup_entry(hass, config_entry): return True -async def async_unload_entry(hass, config_entry): +async def async_unload_entry( + hass: HomeAssistant, config_entry: config_entries.ConfigEntry +): """Unload ZHA config entry.""" await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].shutdown() await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_update_device_storage() @@ -139,10 +142,6 @@ async def async_unload_entry(hass, config_entry): GROUP_PROBE.cleanup() api.async_unload_api(hass) - dispatchers = hass.data[DATA_ZHA].get(DATA_ZHA_DISPATCHERS, []) - for unsub_dispatcher in dispatchers: - unsub_dispatcher() - # our components don't have unload methods so no need to look at return values await asyncio.gather( *( diff --git a/homeassistant/components/zha/alarm_control_panel.py b/homeassistant/components/zha/alarm_control_panel.py index c76e26baaf6..3e83a5cd3d1 100644 --- a/homeassistant/components/zha/alarm_control_panel.py +++ b/homeassistant/components/zha/alarm_control_panel.py @@ -12,6 +12,7 @@ from homeassistant.components.alarm_control_panel import ( AlarmControlPanelEntity, ) from homeassistant.components.zha.core.typing import ZhaDeviceType +from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( STATE_ALARM_ARMED_AWAY, STATE_ALARM_ARMED_HOME, @@ -20,8 +21,9 @@ from homeassistant.const import ( STATE_ALARM_TRIGGERED, Platform, ) -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback from .core import discovery from .core.channels.security import ( @@ -35,7 +37,6 @@ from .core.const import ( CONF_ALARM_FAILED_TRIES, CONF_ALARM_MASTER_CODE, DATA_ZHA, - DATA_ZHA_DISPATCHERS, SIGNAL_ADD_ENTITIES, ZHA_ALARM_OPTIONS, ) @@ -56,7 +57,11 @@ IAS_ACE_STATE_MAP = { } -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation alarm control panel from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.ALARM_CONTROL_PANEL] @@ -67,7 +72,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): discovery.async_add_entities, async_add_entities, entities_to_create ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) @STRICT_MATCH(channel_names=CHANNEL_IAS_ACE) diff --git a/homeassistant/components/zha/binary_sensor.py b/homeassistant/components/zha/binary_sensor.py index 7e083a07d72..c0af9a807d5 100644 --- a/homeassistant/components/zha/binary_sensor.py +++ b/homeassistant/components/zha/binary_sensor.py @@ -12,9 +12,11 @@ from homeassistant.components.binary_sensor import ( DEVICE_CLASS_VIBRATION, BinarySensorEntity, ) +from homeassistant.config_entries import ConfigEntry from homeassistant.const import STATE_ON, Platform -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback from .core import discovery from .core.const import ( @@ -24,7 +26,6 @@ from .core.const import ( CHANNEL_ON_OFF, CHANNEL_ZONE, DATA_ZHA, - DATA_ZHA_DISPATCHERS, SIGNAL_ADD_ENTITIES, SIGNAL_ATTR_UPDATED, ) @@ -44,7 +45,11 @@ CLASS_MAPPING = { STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, Platform.BINARY_SENSOR) -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation binary sensor from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.BINARY_SENSOR] @@ -55,7 +60,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): discovery.async_add_entities, async_add_entities, entities_to_create ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) class BinarySensor(ZhaEntity, BinarySensorEntity): diff --git a/homeassistant/components/zha/climate.py b/homeassistant/components/zha/climate.py index 7cb2c0f16e4..9c01f6630db 100644 --- a/homeassistant/components/zha/climate.py +++ b/homeassistant/components/zha/climate.py @@ -39,14 +39,16 @@ from homeassistant.components.climate.const import ( SUPPORT_TARGET_TEMPERATURE, SUPPORT_TARGET_TEMPERATURE_RANGE, ) +from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( ATTR_TEMPERATURE, PRECISION_TENTHS, TEMP_CELSIUS, Platform, ) -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.event import async_track_time_interval import homeassistant.util.dt as dt_util @@ -55,7 +57,6 @@ from .core.const import ( CHANNEL_FAN, CHANNEL_THERMOSTAT, DATA_ZHA, - DATA_ZHA_DISPATCHERS, PRESET_COMPLEX, PRESET_SCHEDULE, SIGNAL_ADD_ENTITIES, @@ -154,7 +155,11 @@ SYSTEM_MODE_2_HVAC = { ZCL_TEMP = 100 -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation sensor from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.CLIMATE] unsub = async_dispatcher_connect( @@ -164,7 +169,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): discovery.async_add_entities, async_add_entities, entities_to_create ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) @MULTI_MATCH(channel_names=CHANNEL_THERMOSTAT, aux_channels=CHANNEL_FAN) diff --git a/homeassistant/components/zha/core/const.py b/homeassistant/components/zha/core/const.py index 876d4c4e9f5..585124c3ee9 100644 --- a/homeassistant/components/zha/core/const.py +++ b/homeassistant/components/zha/core/const.py @@ -165,7 +165,6 @@ DATA_ZHA = "zha" DATA_ZHA_CONFIG = "config" DATA_ZHA_BRIDGE_ID = "zha_bridge_id" DATA_ZHA_CORE_EVENTS = "zha_core_events" -DATA_ZHA_DISPATCHERS = "zha_dispatchers" DATA_ZHA_GATEWAY = "zha_gateway" DATA_ZHA_PLATFORM_LOADED = "platform_loaded" DATA_ZHA_SHUTDOWN_TASK = "zha_shutdown_task" diff --git a/homeassistant/components/zha/cover.py b/homeassistant/components/zha/cover.py index 2febc209755..8b9bf5589b3 100644 --- a/homeassistant/components/zha/cover.py +++ b/homeassistant/components/zha/cover.py @@ -14,6 +14,7 @@ from homeassistant.components.cover import ( DEVICE_CLASS_SHADE, CoverEntity, ) +from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( STATE_CLOSED, STATE_CLOSING, @@ -21,8 +22,9 @@ from homeassistant.const import ( STATE_OPENING, Platform, ) -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback from .core import discovery from .core.const import ( @@ -31,7 +33,6 @@ from .core.const import ( CHANNEL_ON_OFF, CHANNEL_SHADE, DATA_ZHA, - DATA_ZHA_DISPATCHERS, SIGNAL_ADD_ENTITIES, SIGNAL_ATTR_UPDATED, SIGNAL_SET_LEVEL, @@ -45,7 +46,11 @@ _LOGGER = logging.getLogger(__name__) STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, Platform.COVER) -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation cover from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.COVER] @@ -56,7 +61,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): discovery.async_add_entities, async_add_entities, entities_to_create ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) @STRICT_MATCH(channel_names=CHANNEL_COVER) diff --git a/homeassistant/components/zha/device_tracker.py b/homeassistant/components/zha/device_tracker.py index c203462d44d..4219bfd6288 100644 --- a/homeassistant/components/zha/device_tracker.py +++ b/homeassistant/components/zha/device_tracker.py @@ -4,15 +4,16 @@ import time from homeassistant.components.device_tracker import SOURCE_TYPE_ROUTER from homeassistant.components.device_tracker.config_entry import ScannerEntity +from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback from .core import discovery from .core.const import ( CHANNEL_POWER_CONFIGURATION, DATA_ZHA, - DATA_ZHA_DISPATCHERS, SIGNAL_ADD_ENTITIES, SIGNAL_ATTR_UPDATED, ) @@ -23,7 +24,11 @@ from .sensor import Battery STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, Platform.DEVICE_TRACKER) -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation device tracker from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.DEVICE_TRACKER] @@ -34,7 +39,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): discovery.async_add_entities, async_add_entities, entities_to_create ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) @STRICT_MATCH(channel_names=CHANNEL_POWER_CONFIGURATION) diff --git a/homeassistant/components/zha/fan.py b/homeassistant/components/zha/fan.py index 0402e93ff8c..729a0f7e618 100644 --- a/homeassistant/components/zha/fan.py +++ b/homeassistant/components/zha/fan.py @@ -15,9 +15,11 @@ from homeassistant.components.fan import ( FanEntity, NotValidPresetModeError, ) +from homeassistant.config_entries import ConfigEntry from homeassistant.const import STATE_UNAVAILABLE, Platform -from homeassistant.core import State, callback +from homeassistant.core import HomeAssistant, State, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.util.percentage import ( int_states_in_range, percentage_to_ranged_value, @@ -25,13 +27,7 @@ from homeassistant.util.percentage import ( ) from .core import discovery -from .core.const import ( - CHANNEL_FAN, - DATA_ZHA, - DATA_ZHA_DISPATCHERS, - SIGNAL_ADD_ENTITIES, - SIGNAL_ATTR_UPDATED, -) +from .core.const import CHANNEL_FAN, DATA_ZHA, SIGNAL_ADD_ENTITIES, SIGNAL_ATTR_UPDATED from .core.registries import ZHA_ENTITIES from .entity import ZhaEntity, ZhaGroupEntity @@ -56,7 +52,11 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, Platform.FAN) GROUP_MATCH = functools.partial(ZHA_ENTITIES.group_match, Platform.FAN) -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation fan from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.FAN] @@ -70,7 +70,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): update_before_add=False, ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) class BaseFan(FanEntity): diff --git a/homeassistant/components/zha/light.py b/homeassistant/components/zha/light.py index 476bb733cc9..2c3b6249cf9 100644 --- a/homeassistant/components/zha/light.py +++ b/homeassistant/components/zha/light.py @@ -30,18 +30,20 @@ from homeassistant.components.light import ( SUPPORT_FLASH, SUPPORT_TRANSITION, ) +from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( ATTR_SUPPORTED_FEATURES, STATE_ON, STATE_UNAVAILABLE, Platform, ) -from homeassistant.core import State, callback +from homeassistant.core import HomeAssistant, State, callback from homeassistant.helpers.debounce import Debouncer from homeassistant.helpers.dispatcher import ( async_dispatcher_connect, async_dispatcher_send, ) +from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.event import async_track_time_interval import homeassistant.util.color as color_util @@ -52,7 +54,6 @@ from .core.const import ( CHANNEL_ON_OFF, CONF_DEFAULT_LIGHT_TRANSITION, DATA_ZHA, - DATA_ZHA_DISPATCHERS, EFFECT_BLINK, EFFECT_BREATHE, EFFECT_DEFAULT_VARIANT, @@ -105,7 +106,11 @@ class LightColorMode(enum.IntEnum): COLOR_TEMP = 0x02 -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation light from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.LIGHT] @@ -116,7 +121,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): discovery.async_add_entities, async_add_entities, entities_to_create ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) class BaseLight(LogMixin, light.LightEntity): diff --git a/homeassistant/components/zha/lock.py b/homeassistant/components/zha/lock.py index 0c8ca756785..61104a8b7e9 100644 --- a/homeassistant/components/zha/lock.py +++ b/homeassistant/components/zha/lock.py @@ -5,8 +5,9 @@ import voluptuous as vol from zigpy.zcl.foundation import Status from homeassistant.components.lock import STATE_LOCKED, STATE_UNLOCKED, LockEntity +from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import config_validation as cv, entity_platform from homeassistant.helpers.dispatcher import async_dispatcher_connect @@ -14,7 +15,6 @@ from .core import discovery from .core.const import ( CHANNEL_DOORLOCK, DATA_ZHA, - DATA_ZHA_DISPATCHERS, SIGNAL_ADD_ENTITIES, SIGNAL_ATTR_UPDATED, ) @@ -33,7 +33,11 @@ SERVICE_DISABLE_LOCK_USER_CODE = "disable_lock_user_code" SERVICE_CLEAR_LOCK_USER_CODE = "clear_lock_user_code" -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: entity_platform.AddEntitiesCallback, +): """Set up the Zigbee Home Automation Door Lock from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.LOCK] @@ -44,7 +48,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): discovery.async_add_entities, async_add_entities, entities_to_create ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) platform = entity_platform.async_get_current_platform() diff --git a/homeassistant/components/zha/number.py b/homeassistant/components/zha/number.py index 0acda49c375..e3de49756c2 100644 --- a/homeassistant/components/zha/number.py +++ b/homeassistant/components/zha/number.py @@ -3,15 +3,16 @@ import functools import logging from homeassistant.components.number import NumberEntity +from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback from .core import discovery from .core.const import ( CHANNEL_ANALOG_OUTPUT, DATA_ZHA, - DATA_ZHA_DISPATCHERS, SIGNAL_ADD_ENTITIES, SIGNAL_ATTR_UPDATED, ) @@ -234,7 +235,11 @@ ICONS = { } -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation Analog Output from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.NUMBER] @@ -248,7 +253,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): update_before_add=False, ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) @STRICT_MATCH(channel_names=CHANNEL_ANALOG_OUTPUT) diff --git a/homeassistant/components/zha/sensor.py b/homeassistant/components/zha/sensor.py index d4c508a8378..f14e8b561c3 100644 --- a/homeassistant/components/zha/sensor.py +++ b/homeassistant/components/zha/sensor.py @@ -72,7 +72,6 @@ from .core.const import ( CHANNEL_TEMPERATURE, CHANNEL_THERMOSTAT, DATA_ZHA, - DATA_ZHA_DISPATCHERS, SIGNAL_ADD_ENTITIES, SIGNAL_ATTR_UPDATED, ) @@ -121,7 +120,7 @@ async def async_setup_entry( update_before_add=False, ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) class Sensor(ZhaEntity, SensorEntity): diff --git a/homeassistant/components/zha/switch.py b/homeassistant/components/zha/switch.py index da5e843c59f..b67ab72fa35 100644 --- a/homeassistant/components/zha/switch.py +++ b/homeassistant/components/zha/switch.py @@ -8,15 +8,16 @@ from zigpy.zcl.clusters.general import OnOff from zigpy.zcl.foundation import Status from homeassistant.components.switch import SwitchEntity +from homeassistant.config_entries import ConfigEntry from homeassistant.const import STATE_ON, STATE_UNAVAILABLE, Platform -from homeassistant.core import State, callback +from homeassistant.core import HomeAssistant, State, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback from .core import discovery from .core.const import ( CHANNEL_ON_OFF, DATA_ZHA, - DATA_ZHA_DISPATCHERS, SIGNAL_ADD_ENTITIES, SIGNAL_ATTR_UPDATED, ) @@ -27,7 +28,11 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, Platform.SWITCH) GROUP_MATCH = functools.partial(ZHA_ENTITIES.group_match, Platform.SWITCH) -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +): """Set up the Zigbee Home Automation switch from config entry.""" entities_to_create = hass.data[DATA_ZHA][Platform.SWITCH] @@ -38,7 +43,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): discovery.async_add_entities, async_add_entities, entities_to_create ), ) - hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub) + config_entry.async_on_unload(unsub) class BaseSwitch(SwitchEntity):