Index config entries by domain (#56316)

This commit is contained in:
J. Nick Koston 2021-09-17 10:39:00 -10:00 committed by GitHub
parent 5b0e00a74b
commit e880f1c8f9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 4 deletions

View File

@ -769,6 +769,7 @@ class ConfigEntries:
self.options = OptionsFlowManager(hass) self.options = OptionsFlowManager(hass)
self._hass_config = hass_config self._hass_config = hass_config
self._entries: dict[str, ConfigEntry] = {} self._entries: dict[str, ConfigEntry] = {}
self._domain_index: dict[str, list[str]] = {}
self._store = hass.helpers.storage.Store(STORAGE_VERSION, STORAGE_KEY) self._store = hass.helpers.storage.Store(STORAGE_VERSION, STORAGE_KEY)
EntityRegistryDisabledHandler(hass).async_setup() EntityRegistryDisabledHandler(hass).async_setup()
@ -796,7 +797,9 @@ class ConfigEntries:
"""Return all entries or entries for a specific domain.""" """Return all entries or entries for a specific domain."""
if domain is None: if domain is None:
return list(self._entries.values()) return list(self._entries.values())
return [entry for entry in self._entries.values() if entry.domain == domain] return [
self._entries[entry_id] for entry_id in self._domain_index.get(domain, [])
]
async def async_add(self, entry: ConfigEntry) -> None: async def async_add(self, entry: ConfigEntry) -> None:
"""Add and setup an entry.""" """Add and setup an entry."""
@ -805,6 +808,7 @@ class ConfigEntries:
f"An entry with the id {entry.entry_id} already exists." f"An entry with the id {entry.entry_id} already exists."
) )
self._entries[entry.entry_id] = entry self._entries[entry.entry_id] = entry
self._domain_index.setdefault(entry.domain, []).append(entry.entry_id)
await self.async_setup(entry.entry_id) await self.async_setup(entry.entry_id)
self._async_schedule_save() self._async_schedule_save()
@ -823,6 +827,9 @@ class ConfigEntries:
await entry.async_remove(self.hass) await entry.async_remove(self.hass)
del self._entries[entry.entry_id] del self._entries[entry.entry_id]
self._domain_index[entry.domain].remove(entry.entry_id)
if not self._domain_index[entry.domain]:
del self._domain_index[entry.domain]
self._async_schedule_save() self._async_schedule_save()
dev_reg, ent_reg = await asyncio.gather( dev_reg, ent_reg = await asyncio.gather(
@ -881,9 +888,11 @@ class ConfigEntries:
if config is None: if config is None:
self._entries = {} self._entries = {}
self._domain_index = {}
return return
entries = {} entries = {}
domain_index: dict[str, list[str]] = {}
for entry in config["entries"]: for entry in config["entries"]:
pref_disable_new_entities = entry.get("pref_disable_new_entities") pref_disable_new_entities = entry.get("pref_disable_new_entities")
@ -894,10 +903,13 @@ class ConfigEntries:
"disable_new_entities" "disable_new_entities"
) )
entries[entry["entry_id"]] = ConfigEntry( domain = entry["domain"]
entry_id = entry["entry_id"]
entries[entry_id] = ConfigEntry(
version=entry["version"], version=entry["version"],
domain=entry["domain"], domain=domain,
entry_id=entry["entry_id"], entry_id=entry_id,
data=entry["data"], data=entry["data"],
source=entry["source"], source=entry["source"],
title=entry["title"], title=entry["title"],
@ -911,7 +923,9 @@ class ConfigEntries:
pref_disable_new_entities=pref_disable_new_entities, pref_disable_new_entities=pref_disable_new_entities,
pref_disable_polling=entry.get("pref_disable_polling"), pref_disable_polling=entry.get("pref_disable_polling"),
) )
domain_index.setdefault(domain, []).append(entry_id)
self._domain_index = domain_index
self._entries = entries self._entries = entries
async def async_setup(self, entry_id: str) -> bool: async def async_setup(self, entry_id: str) -> bool:

View File

@ -771,10 +771,14 @@ class MockConfigEntry(config_entries.ConfigEntry):
def add_to_hass(self, hass): def add_to_hass(self, hass):
"""Test helper to add entry to hass.""" """Test helper to add entry to hass."""
hass.config_entries._entries[self.entry_id] = self hass.config_entries._entries[self.entry_id] = self
hass.config_entries._domain_index.setdefault(self.domain, []).append(
self.entry_id
)
def add_to_manager(self, manager): def add_to_manager(self, manager):
"""Test helper to add entry to entry manager.""" """Test helper to add entry to entry manager."""
manager._entries[self.entry_id] = self manager._entries[self.entry_id] = self
manager._domain_index.setdefault(self.domain, []).append(self.entry_id)
def patch_yaml_files(files_dict, endswith=True): def patch_yaml_files(files_dict, endswith=True):