Clean up device update, add via-hub

homeassistant/helpers/device_registry.py

@@ -10,6 +10,7 @@
 from homeassistant.loader import bind_hass
 
 _LOGGER = logging.getLogger(__name__)
+_UNDEF = object()
 
 DATA_REGISTRY = 'device_registry'
 
@@ -32,6 +33,7 @@ class DeviceEntry:
     model = attr.ib(type=str)
     name = attr.ib(type=str, default=None)
     sw_version = attr.ib(type=str, default=None)
+    hub_device_id = attr.ib(type=str, default=None)
     id = attr.ib(type=str, default=attr.Factory(lambda: uuid.uuid4().hex))
 
 
@@ -54,53 +56,99 @@ def async_get_device(self, identifiers: set, connections: set):
         return None
 
     @callback
-    def async_get_or_create(self, *, config_entry, connections, identifiers,
-                            manufacturer, model, name=None, sw_version=None):
+    def async_get_or_create(self, *, config_entry_id, connections, identifiers,
+                            manufacturer, model, name=None, sw_version=None,
+                            via_hub=None):
         """Get device. Create if it doesn't exist."""
         if not identifiers and not connections:
             return None
 
         device = self.async_get_device(identifiers, connections)
 
+        if via_hub is not None:
+            hub_device = self.async_get_device({via_hub}, set())
+            hub_device_id = hub_device.id if hub_device else None
+        else:
+            hub_device_id = None
+
         if device is not None:
-            if config_entry not in device.config_entries:
-                device.config_entries.add(config_entry)
-                self.async_schedule_save()
-            return device
+            return self._async_update_device(
+                device.id, config_entry_id=config_entry_id,
+                hub_device_id=hub_device_id
+            )
 
         device = DeviceEntry(
-            config_entries=[config_entry],
+            config_entries={config_entry_id},
             connections=connections,
             identifiers=identifiers,
             manufacturer=manufacturer,
             model=model,
             name=name,
-            sw_version=sw_version
+            sw_version=sw_version,
+            hub_device_id=hub_device_id
         )
         self.devices[device.id] = device
 
         self.async_schedule_save()
 
         return device
 
+    @callback
+    def _async_update_device(self, device_id, *, config_entry_id=_UNDEF,
+                             remove_config_entry_id=_UNDEF,
+                             hub_device_id=_UNDEF):
+        """Update device attributes."""
+        old = self.devices[device_id]
+
+        changes = {}
+
+        config_entries = old.config_entries
+
+        if (config_entry_id is not _UNDEF and
+                config_entry_id not in old.config_entries):
+            config_entries = old.config_entries | {config_entry_id}
+
+        if (remove_config_entry_id is not _UNDEF and
+                remove_config_entry_id in config_entries):
+            config_entries = set(config_entries)
+            config_entries.remove(remove_config_entry_id)
+
+        if config_entries is not old.config_entries:
+            changes['config_entries'] = config_entries
+
+        if (hub_device_id is not _UNDEF and
+                hub_device_id != old.hub_device_id):
+            changes['hub_device_id'] = hub_device_id
+
+        if not changes:
+            return old
+
+        new = self.devices[device_id] = attr.evolve(old, **changes)
+        self.async_schedule_save()
+        return new
+
     async def async_load(self):
         """Load the device registry."""
-        devices = await self._store.async_load()
-
-        if devices is None:
-            self.devices = OrderedDict()
-            return
-
-        self.devices = {device['id']: DeviceEntry(
-            config_entries=device['config_entries'],
-            connections={tuple(conn) for conn in device['connections']},
-            identifiers={tuple(iden) for iden in device['identifiers']},
-            manufacturer=device['manufacturer'],
-            model=device['model'],
-            name=device['name'],
-            sw_version=device['sw_version'],
-            id=device['id'],
-        ) for device in devices['devices']}
+        data = await self._store.async_load()
+
+        devices = OrderedDict()
+
+        if data is not None:
+            for device in data['devices']:
+                devices[device['id']] = DeviceEntry(
+                    config_entries=set(device['config_entries']),
+                    connections={tuple(conn) for conn in device['connections']},
+                    identifiers={tuple(iden) for iden in device['identifiers']},
+                    manufacturer=device['manufacturer'],
+                    model=device['model'],
+                    name=device['name'],
+                    sw_version=device['sw_version'],
+                    id=device['id'],
+                    # Introduced in 0.79
+                    hub_device_id=device.get('hub_device_id'),
+                )
+
+        self.devices = devices
 
     @callback
     def async_schedule_save(self):
@@ -122,18 +170,19 @@ def _data_to_save(self):
                 'name': entry.name,
                 'sw_version': entry.sw_version,
                 'id': entry.id,
+                'hub_device_id': entry.hub_device_id,
             } for entry in self.devices.values()
         ]
 
         return data
 
     @callback
-    def async_clear_config_entry(self, config_entry):
+    def async_clear_config_entry(self, config_entry_id):
         """Clear config entry from registry entries."""
-        for device in self.devices.values():
-            if config_entry in device.config_entries:
-                device.config_entries.remove(config_entry)
-                self.async_schedule_save()
+        for dev_id, device in self.devices.items():
+            if config_entry_id in device.config_entries:
+                self._async_update_device(
+                    dev_id, remove_config_entry_id=config_entry_id)
 
 
 @bind_hass

homeassistant/helpers/entity_platform.py

@@ -273,16 +273,19 @@ async def _async_add_entity(self, entity, update_before_add,
                 config_entry_id = None
 
             device_info = entity.device_info
+
             if config_entry_id is not None and device_info is not None:
                 device = device_registry.async_get_or_create(
-                    config_entry=config_entry_id,
-                    connections=device_info.get('connections', []),
-                    identifiers=device_info.get('identifiers', []),
+                    config_entry_id=config_entry_id,
+                    connections=device_info.get('connections') or set(),
+                    identifiers=device_info.get('identifiers') or set(),
                     manufacturer=device_info.get('manufacturer'),
                     model=device_info.get('model'),
                     name=device_info.get('name'),
-                    sw_version=device_info.get('sw_version'))
-                device_id = device.id
+                    sw_version=device_info.get('sw_version'),
+                    via_hub=device_info.get('via_hub'))
+                if device:
+                    device_id = device.id
             else:
                 device_id = None
 

homeassistant/helpers/entity_registry.py

@@ -31,7 +31,7 @@
 STORAGE_KEY = 'core.entity_registry'
 
 
-@attr.s(slots=True)
+@attr.s(slots=True, frozen=True)
 class RegistryEntry:
     """Entity Registry Entry."""
 
@@ -113,14 +113,9 @@ def async_get_or_create(self, domain, platform, unique_id, *,
         """Get entity. Create if it doesn't exist."""
         entity_id = self.async_get_entity_id(domain, platform, unique_id)
         if entity_id:
-            entry = self.entities[entity_id]
-            if entry.config_entry_id == config_entry_id:
-                return entry
-
-            self._async_update_entity(
+            return self._async_update_entity(
                 entity_id, config_entry_id=config_entry_id,
                 device_id=device_id)
-            return self.entities[entity_id]
 
         entity_id = self.async_generate_entity_id(
             domain, suggested_object_id or '{}_{}'.format(platform, unique_id))
@@ -253,10 +248,9 @@ def _data_to_save(self):
     @callback
     def async_clear_config_entry(self, config_entry):
         """Clear config entry from registry entries."""
-        for entry in self.entities.values():
+        for entity_id, entry in self.entities.items():
             if config_entry == entry.config_entry_id:
-                entry.config_entry_id = None
-                self.async_schedule_save()
+                self._async_update_entity(entity_id, config_entry_id=None)
 
 
 @bind_hass

tests/common.py

@@ -763,6 +763,11 @@ def available(self):
         """Return True if entity is available."""
         return self._handle('available')
 
+    @property
+    def device_info(self):
+        """Info how it links to a device."""
+        return self._handle('device_info')
+
     def _handle(self, attr):
         """Return attribute value."""
         if attr in self._values: