☁️ improve get cloud devices

custom_components/xiaomi_miot/__init__.py

@@ -276,16 +276,15 @@ async def async_setup_xiaomi_cloud(hass: hass_core.HomeAssistant, config_entry:
     try:
         cloud = await entry.get_cloud(check=True)
         config[CONF_XIAOMI_CLOUD] = cloud
-        config['devices_by_mac'] = await cloud.async_get_devices_by_key('mac', filters=entry_config) or {}
+        devices = await entry.get_cloud_devices()
     except (MiCloudException, MiCloudAccessDenied) as exc:
         _LOGGER.error('Setup xiaomi cloud for user: %s failed: %s', username, exc)
         return False
-    if not config.get('devices_by_mac'):
+    if not devices:
         _LOGGER.warning('None device in xiaomi cloud: %s', username)
     else:
-        cnt = len(config['devices_by_mac'])
-        _LOGGER.debug('Setup xiaomi cloud for user: %s, %s devices', username, cnt)
-    for mac, d in config['devices_by_mac'].items():
+        _LOGGER.debug('Setup xiaomi cloud for user: %s, %s devices', username, len(devices))
+    for d in devices.values():
         device = await entry.new_device(d)
         if not device.spec:
             _LOGGER.warning('%s: Device has no spec %s', device.name_model, device.info.urn)
@@ -577,18 +576,19 @@ async def updated(event: hass_core.Event):
     hass.bus.async_listen(dr.EVENT_DEVICE_REGISTRY_UPDATED, updated)
 
 
-async def async_remove_config_entry_device(hass: hass_core.HomeAssistant, entry: ConfigEntry, device: dr.DeviceEntry):
+async def async_remove_config_entry_device(hass: hass_core.HomeAssistant, config_entry: ConfigEntry, device: dr.DeviceEntry):
     """Supported from Hass v2022.3"""
-    dvc = None
+    entry = HassEntry.init(hass, config_entry)
+    cloud_device = None
     identifier = next(iter(device.identifiers))
     if len(identifier) >= 2 and identifier[0] == DOMAIN:
         mac = identifier[1].split('-')[0]
         if mac:
-            dvc = hass.data[DOMAIN].get(entry.entry_id, {}).get('devices_by_mac', {}).get(mac.upper()) or {}
-    data = {**entry.data, **entry.options}
-    for typ in (['did'] if dvc else []):
+            cloud_device = await entry.get_cloud_device(mac=mac.upper())
+    data = entry.entry.data
+    for typ in (['did'] if cloud_device else []):
         filter_typ = data.get(f'filter_{typ}')
-        filter_val = dvc.get(typ)
+        filter_val = cloud_device.get(typ)
         if not filter_val or not filter_typ:
             continue
         lst = data.get(f'{typ}_list') or []
@@ -597,8 +597,8 @@ async def async_remove_config_entry_device(hass: hass_core.HomeAssistant, entry:
         else:
             lst = list({*lst}.difference({filter_val}))
         data[f'{typ}_list'] = lst
-        hass.config_entries.async_update_entry(entry, data=data)
-        _LOGGER.info('Remove miot device: %s', [dvc])
+        hass.config_entries.async_update_entry(config_entry, data=data)
+        _LOGGER.info('Remove miot device: %s', cloud_device)
 
     dr.async_get(hass).async_remove_device(device.id)
     return True
@@ -1007,12 +1007,9 @@ def miot_device(self):
     @property
     def miot_did(self):
         did = self.custom_config('miot_did') or self._config.get('miot_did')
-        if self.entity_id and not did:
-            mac = self.device.info.mac
-            dvs = self.entry_config('devices_by_mac') or {}
-            if mac in dvs:
-                return dvs[mac].get('did')
-        return did
+        if did:
+            return did
+        return self.device.did
 
     @property
     def xiaomi_cloud(self):

custom_components/xiaomi_miot/core/hass_entry.py

@@ -16,13 +16,15 @@
 class HassEntry:
     ALL: dict[str, 'HassEntry'] = {}
     cloud: MiotCloud = None
+    cloud_devices = None
 
     def __init__(self, hass: HomeAssistant, entry: ConfigEntry):
         self.id = entry.entry_id
         self.hass = hass
         self.entry = entry
         self.adders: dict[str, AddEntitiesCallback] = {}
         self.devices: dict[str, 'Device'] = {}
+        self.mac_to_did = {}
 
     @staticmethod
     def init(hass: HomeAssistant, entry: ConfigEntry):
@@ -106,3 +108,22 @@ async def get_cloud(self, check=False, login=False):
         if check:
             await self.cloud.async_check_auth(notify=True)
         return self.cloud
+
+    async def get_cloud_devices(self):
+        if isinstance(self.cloud_devices, dict):
+            return self.cloud_devices
+        cloud = await self.get_cloud()
+        config = self.get_config()
+        self.cloud_devices = await cloud.async_get_devices_by_key('did', filters=config) or {}
+        for did, info in self.cloud_devices.items():
+            mac = info.get('mac') or did
+            self.mac_to_did[mac] = did
+        return self.cloud_devices
+
+    async def get_cloud_device(self, did=None, mac=None):
+        devices = await self.get_cloud_devices()
+        if mac and not did:
+            did = self.mac_to_did.get(mac)
+        if did:
+            return devices.get(did)
+        return None