diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py index d69cdd3d997..762e6813b1d 100644 --- a/homeassistant/helpers/entity.py +++ b/homeassistant/helpers/entity.py @@ -10,6 +10,7 @@ from homeassistant.const import ( ATTR_UNIT_OF_MEASUREMENT, DEVICE_DEFAULT_NAME, STATE_OFF, STATE_ON, STATE_UNAVAILABLE, STATE_UNKNOWN, TEMP_CELSIUS, TEMP_FAHRENHEIT, ATTR_ENTITY_PICTURE, ATTR_SUPPORTED_FEATURES, ATTR_DEVICE_CLASS) +from homeassistant.helpers.entity_registry import EVENT_ENTITY_REGISTRY_UPDATED from homeassistant.core import HomeAssistant, callback from homeassistant.config import DATA_CUSTOMIZE from homeassistant.exceptions import NoEntitySpecifiedError @@ -78,8 +79,8 @@ class Entity: # Process updates in parallel parallel_updates = None - # Name in the entity registry - registry_name = None + # Entry in the entity registry + registry_entry = None # Hold list for functions to call on remove. _on_remove = None @@ -259,7 +260,9 @@ class Entity: if unit_of_measurement is not None: attr[ATTR_UNIT_OF_MEASUREMENT] = unit_of_measurement - name = self.registry_name or self.name + entry = self.registry_entry + # pylint: disable=consider-using-ternary + name = (entry and entry.name) or self.name if name is not None: attr[ATTR_FRIENDLY_NAME] = name @@ -391,6 +394,7 @@ class Entity: async def async_remove(self): """Remove entity from Home Assistant.""" + await self.async_internal_will_remove_from_hass() await self.async_will_remove_from_hass() if self._on_remove is not None: @@ -399,27 +403,52 @@ class Entity: self.hass.states.async_remove(self.entity_id) - @callback - def async_registry_updated(self, old, new): - """Handle entity registry update.""" - self.registry_name = new.name - - if new.entity_id == self.entity_id: - self.async_schedule_update_ha_state() - return - - async def readd(): - """Remove and add entity again.""" - await self.async_remove() - await self.platform.async_add_entities([self]) - - self.hass.async_create_task(readd()) - async def async_added_to_hass(self) -> None: - """Run when entity about to be added to hass.""" + """Run when entity about to be added to hass. + + To be extended by integrations. + """ async def async_will_remove_from_hass(self) -> None: - """Run when entity will be removed from hass.""" + """Run when entity will be removed from hass. + + To be extended by integrations. + """ + + async def async_internal_added_to_hass(self) -> None: + """Run when entity about to be added to hass. + + Not to be extended by integrations. + """ + if self.registry_entry is not None: + self.async_on_remove(self.hass.bus.async_listen( + EVENT_ENTITY_REGISTRY_UPDATED, self._async_registry_updated)) + + async def async_internal_will_remove_from_hass(self) -> None: + """Run when entity will be removed from hass. + + Not to be extended by integrations. + """ + + async def _async_registry_updated(self, event): + """Handle entity registry update.""" + data = event.data + if data['action'] != 'update' and data.get( + 'old_entity_id', data['entity_id']) != self.entity_id: + return + + ent_reg = await self.hass.helpers.entity_registry.async_get_registry() + old = self.registry_entry + self.registry_entry = ent_reg.async_get(data['entity_id']) + + if self.registry_entry.entity_id == old.entity_id: + self.async_write_ha_state() + return + + await self.async_remove() + + self.entity_id = self.registry_entry.entity_id + await self.platform.async_add_entities([self]) def __eq__(self, other): """Return the comparison.""" diff --git a/homeassistant/helpers/entity_platform.py b/homeassistant/helpers/entity_platform.py index 8b1b8502586..12e0e01acaf 100644 --- a/homeassistant/helpers/entity_platform.py +++ b/homeassistant/helpers/entity_platform.py @@ -320,9 +320,8 @@ class EntityPlatform: '"{} {}"'.format(self.platform_name, entity.unique_id)) return + entity.registry_entry = entry entity.entity_id = entry.entity_id - entity.registry_name = entry.name - entity.async_on_remove(entry.add_update_listener(entity)) # We won't generate an entity ID if the platform has already set one # We will however make sure that platform cannot pick a registered ID @@ -360,6 +359,7 @@ class EntityPlatform: self.entities[entity_id] = entity entity.async_on_remove(lambda: self.entities.pop(entity_id)) + await entity.async_internal_added_to_hass() await entity.async_added_to_hass() await entity.async_update_ha_state() diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 302eda79426..6d3a8a42655 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -12,7 +12,6 @@ from collections import OrderedDict from itertools import chain import logging from typing import List, Optional, cast -import weakref import attr @@ -50,8 +49,6 @@ class RegistryEntry: disabled_by = attr.ib( type=str, default=None, validator=attr.validators.in_((DISABLED_HASS, DISABLED_USER, None))) - update_listeners = attr.ib(type=list, default=attr.Factory(list), - repr=False) domain = attr.ib(type=str, init=False, repr=False) @domain.default @@ -64,18 +61,6 @@ class RegistryEntry: """Return if entry is disabled.""" return self.disabled_by is not None - def add_update_listener(self, listener): - """Listen for when entry is updated. - - Listener: Callback function(old_entry, new_entry) - - Returns function to unlisten. - """ - weak_listener = weakref.ref(listener) - self.update_listeners.append(weak_listener) - - return lambda: self.update_listeners.remove(weak_listener) - class EntityRegistry: """Class to hold a registry of entities.""" @@ -247,26 +232,17 @@ class EntityRegistry: new = self.entities[entity_id] = attr.evolve(old, **changes) - to_remove = [] - for listener_ref in new.update_listeners: - listener = listener_ref() - if listener is None: - to_remove.append(listener_ref) - else: - try: - listener.async_registry_updated(old, new) - except Exception: # pylint: disable=broad-except - _LOGGER.exception('Error calling update listener') - - for ref in to_remove: - new.update_listeners.remove(ref) - self.async_schedule_save() - self.hass.bus.async_fire(EVENT_ENTITY_REGISTRY_UPDATED, { + data = { 'action': 'update', - 'entity_id': entity_id - }) + 'entity_id': entity_id, + } + + if old.entity_id != entity_id: + data['old_entity_id'] = old.entity_id + + self.hass.bus.async_fire(EVENT_ENTITY_REGISTRY_UPDATED, data) return new diff --git a/homeassistant/helpers/restore_state.py b/homeassistant/helpers/restore_state.py index 355555ec9dc..291bc6d1a0a 100644 --- a/homeassistant/helpers/restore_state.py +++ b/homeassistant/helpers/restore_state.py @@ -186,18 +186,18 @@ class RestoreStateData(): class RestoreEntity(Entity): """Mixin class for restoring previous entity state.""" - async def async_added_to_hass(self) -> None: + async def async_internal_added_to_hass(self) -> None: """Register this entity as a restorable entity.""" _, data = await asyncio.gather( - super().async_added_to_hass(), + super().async_internal_added_to_hass(), RestoreStateData.async_get_instance(self.hass), ) data.async_restore_entity_added(self.entity_id) - async def async_will_remove_from_hass(self) -> None: + async def async_internal_will_remove_from_hass(self) -> None: """Run when entity will be removed from hass.""" _, data = await asyncio.gather( - super().async_will_remove_from_hass(), + super().async_internal_will_remove_from_hass(), RestoreStateData.async_get_instance(self.hass), ) data.async_restore_entity_removed(self.entity_id) diff --git a/tests/helpers/test_restore_state.py b/tests/helpers/test_restore_state.py index bc2ab6937c3..ff38e8fb763 100644 --- a/tests/helpers/test_restore_state.py +++ b/tests/helpers/test_restore_state.py @@ -104,12 +104,12 @@ async def test_dump_data(hass): entity = Entity() entity.hass = hass entity.entity_id = 'input_boolean.b0' - await entity.async_added_to_hass() + await entity.async_internal_added_to_hass() entity = RestoreEntity() entity.hass = hass entity.entity_id = 'input_boolean.b1' - await entity.async_added_to_hass() + await entity.async_internal_added_to_hass() data = await RestoreStateData.async_get_instance(hass) now = dt_util.utcnow() @@ -144,7 +144,7 @@ async def test_dump_data(hass): assert written_states[1]['state']['state'] == 'off' # Test that removed entities are not persisted - await entity.async_will_remove_from_hass() + await entity.async_remove() with patch('homeassistant.helpers.restore_state.Store.async_save' ) as mock_write_data, patch.object( @@ -170,12 +170,12 @@ async def test_dump_error(hass): entity = Entity() entity.hass = hass entity.entity_id = 'input_boolean.b0' - await entity.async_added_to_hass() + await entity.async_internal_added_to_hass() entity = RestoreEntity() entity.hass = hass entity.entity_id = 'input_boolean.b1' - await entity.async_added_to_hass() + await entity.async_internal_added_to_hass() data = await RestoreStateData.async_get_instance(hass) @@ -206,7 +206,7 @@ async def test_state_saved_on_remove(hass): entity = RestoreEntity() entity.hass = hass entity.entity_id = 'input_boolean.b0' - await entity.async_added_to_hass() + await entity.async_internal_added_to_hass() hass.states.async_set('input_boolean.b0', 'on') @@ -215,7 +215,7 @@ async def test_state_saved_on_remove(hass): # No last states should currently be saved assert not data.last_states - await entity.async_will_remove_from_hass() + await entity.async_remove() # We should store the input boolean state when it is removed assert data.last_states['input_boolean.b0'].state.state == 'on'