jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-07-23 13:17:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Speed up initializing config flows (#124015)

Browse Source
This commit is contained in:
J. Nick Koston 2024-08-16 16:48:03 -05:00 committed by GitHub
parent 24680b731f
commit 91951ed734
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 28 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
homeassistant/config_entries.py
View File

@ -3,7 +3,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
from collections import UserDict from collections import UserDict, defaultdict
from collections.abc import ( from collections.abc import (
Callable, Callable,
Coroutine, Coroutine,
@ -1224,8 +1224,12 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
super().__init__(hass) super().__init__(hass)
self.config_entries = config_entries self.config_entries = config_entries
self._hass_config = hass_config self._hass_config = hass_config
self._pending_import_flows: dict[str, dict[str, asyncio.Future[None]]] = {} self._pending_import_flows: defaultdict[
self._initialize_futures: dict[str, list[asyncio.Future[None]]] = {} str, dict[str, asyncio.Future[None]]
] = defaultdict(dict)
self._initialize_futures: defaultdict[str, set[asyncio.Future[None]]] = (
defaultdict(set)
)
self._discovery_debouncer = Debouncer[None]( self._discovery_debouncer = Debouncer[None](
hass, hass,
_LOGGER, _LOGGER,
@ -1278,12 +1282,11 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
loop = self.hass.loop loop = self.hass.loop
if context["source"] == SOURCE_IMPORT: if context["source"] == SOURCE_IMPORT:
self._pending_import_flows.setdefault(handler, {})[flow_id] = ( self._pending_import_flows[handler][flow_id] = loop.create_future()
loop.create_future()
)
cancel_init_future = loop.create_future() cancel_init_future = loop.create_future()
self._initialize_futures.setdefault(handler, []).append(cancel_init_future) handler_init_futures = self._initialize_futures[handler]
handler_init_futures.add(cancel_init_future)
try: try:
async with interrupt( async with interrupt(
cancel_init_future, cancel_init_future,
@ -1294,8 +1297,13 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
except FlowCancelledError as ex: except FlowCancelledError as ex:
raise asyncio.CancelledError from ex raise asyncio.CancelledError from ex
finally: finally:
self._initialize_futures[handler].remove(cancel_init_future) handler_init_futures.remove(cancel_init_future)
self._pending_import_flows.get(handler, {}).pop(flow_id, None) if not handler_init_futures:
del self._initialize_futures[handler]
if handler in self._pending_import_flows:
self._pending_import_flows[handler].pop(flow_id, None)
if not self._pending_import_flows[handler]:
del self._pending_import_flows[handler]
if result["type"] != data_entry_flow.FlowResultType.ABORT: if result["type"] != data_entry_flow.FlowResultType.ABORT:
await self.async_post_init(flow, result) await self.async_post_init(flow, result)
@ -1322,11 +1330,18 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
try: try:
result = await self._async_handle_step(flow, flow.init_step, data) result = await self._async_handle_step(flow, flow.init_step, data)
finally: finally:
init_done = self._pending_import_flows.get(handler, {}).get(flow_id) self._set_pending_import_done(flow)
if init_done and not init_done.done():
init_done.set_result(None)
return flow, result return flow, result
def _set_pending_import_done(self, flow: ConfigFlow) -> None:
"""Set pending import flow as done."""
if (
(handler_import_flows := self._pending_import_flows.get(flow.handler))
and (init_done := handler_import_flows.get(flow.flow_id))
and not init_done.done()
):
init_done.set_result(None)
@callback @callback
def async_shutdown(self) -> None: def async_shutdown(self) -> None:
"""Cancel any initializing flows.""" """Cancel any initializing flows."""
@ -1347,9 +1362,7 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
# We do this to avoid a circular dependency where async_finish_flow sets up a # We do this to avoid a circular dependency where async_finish_flow sets up a
# new entry, which needs the integration to be set up, which is waiting for # new entry, which needs the integration to be set up, which is waiting for
# init to be done. # init to be done.
init_done = self._pending_import_flows.get(flow.handler, {}).get(flow.flow_id) self._set_pending_import_done(flow)
if init_done and not init_done.done():
init_done.set_result(None)
# Remove notification if no other discovery config entries in progress # Remove notification if no other discovery config entries in progress
if not self._async_has_other_discovery_flows(flow.flow_id): if not self._async_has_other_discovery_flows(flow.flow_id):