From d7e461a537379726a473538ac70a4f728357c3b8 Mon Sep 17 00:00:00 2001 From: Marc Mueller <30130371+cdce8p@users.noreply.github.com> Date: Wed, 27 Mar 2024 23:30:01 +0100 Subject: [PATCH] Use HassKey for setup keys --- homeassistant/bootstrap.py | 3 +-- homeassistant/config_entries.py | 4 +-- homeassistant/setup.py | 46 ++++++++++++++------------------- tests/test_setup.py | 7 ----- 4 files changed, 22 insertions(+), 38 deletions(-) diff --git a/homeassistant/bootstrap.py b/homeassistant/bootstrap.py index 03c0de1ff6224..e2728c1ed4b45 100644 --- a/homeassistant/bootstrap.py +++ b/homeassistant/bootstrap.py @@ -881,8 +881,7 @@ async def _async_set_up_integrations( hass: core.HomeAssistant, config: dict[str, Any] ) -> None: """Set up all the integrations.""" - setup_started: dict[tuple[str, str | None], float] = {} - hass.data[DATA_SETUP_STARTED] = setup_started + setup_started = hass.data.setdefault(DATA_SETUP_STARTED, {}) watcher = _WatchPendingSetups(hass, setup_started) watcher.async_start() diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py index 42194641f7f6e..7651641339733 100644 --- a/homeassistant/config_entries.py +++ b/homeassistant/config_entries.py @@ -1953,9 +1953,7 @@ async def async_wait_component(self, entry: ConfigEntry) -> bool: Config entries which are created after Home Assistant is started can't be waited for, the function will just return if the config entry is loaded or not. """ - setup_done: dict[str, asyncio.Future[bool]] = self.hass.data.get( - DATA_SETUP_DONE, {} - ) + setup_done = self.hass.data.get(DATA_SETUP_DONE, {}) if setup_future := setup_done.get(entry.domain): await setup_future # The component was not loaded. diff --git a/homeassistant/setup.py b/homeassistant/setup.py index 178ee6425e322..279df71f71ea4 100644 --- a/homeassistant/setup.py +++ b/homeassistant/setup.py @@ -32,6 +32,7 @@ from .helpers.issue_registry import IssueSeverity, async_create_issue from .helpers.typing import ConfigType from .util.async_ import create_eager_task +from .util.hass_dict import HassKey current_setup_group: contextvars.ContextVar[tuple[str, str | None] | None] = ( contextvars.ContextVar("current_setup_group", default=None) @@ -44,29 +45,30 @@ BASE_PLATFORMS = {platform.value for platform in Platform} -# DATA_SETUP is a dict[str, asyncio.Future[bool]], indicating domains which are currently +# DATA_SETUP is a dict, indicating domains which are currently # being setup or which failed to setup: # - Tasks are added to DATA_SETUP by `async_setup_component`, the key is the domain # being setup and the Task is the `_async_setup_component` helper. # - Tasks are removed from DATA_SETUP if setup was successful, that is, # the task returned True. -DATA_SETUP = "setup_tasks" +DATA_SETUP: HassKey[dict[str, asyncio.Future[bool]]] = HassKey() -# DATA_SETUP_DONE is a dict [str, asyncio.Future[bool]], indicating components which -# will be setup: +# DATA_SETUP_DONE is a dict, indicating components which will be setup: # - Events are added to DATA_SETUP_DONE during bootstrap by # async_set_domains_to_be_loaded, the key is the domain which will be loaded. # - Events are set and removed from DATA_SETUP_DONE when async_setup_component # is finished, regardless of if the setup was successful or not. -DATA_SETUP_DONE = "setup_done" +DATA_SETUP_DONE: HassKey[dict[str, asyncio.Future[bool]]] = HassKey() -# DATA_SETUP_STARTED is a dict [tuple[str, str | None], float], indicating when an attempt +# DATA_SETUP_STARTED is a dict, indicating when an attempt # to setup a component started. -DATA_SETUP_STARTED = "setup_started" +DATA_SETUP_STARTED: HassKey[dict[tuple[str, str | None], float]] = HassKey() -# DATA_SETUP_TIME is a defaultdict[str, defaultdict[str | None, defaultdict[SetupPhases, float]]] -# indicating how time was spent setting up a component and each group (config entry). -DATA_SETUP_TIME = "setup_time" +# DATA_SETUP_TIME is a defaultdict, indicating how time was spent +# setting up a component. +DATA_SETUP_TIME: HassKey[ + defaultdict[str, defaultdict[str | None, defaultdict[SetupPhases, float]]] +] = HassKey() DATA_DEPS_REQS = "deps_reqs_processed" @@ -125,9 +127,7 @@ def async_set_domains_to_be_loaded(hass: core.HomeAssistant, domains: set[str]) - Properly handle after_dependencies. - Keep track of domains which will load but have not yet finished loading """ - setup_done_futures: dict[str, asyncio.Future[bool]] = hass.data.setdefault( - DATA_SETUP_DONE, {} - ) + setup_done_futures = hass.data.setdefault(DATA_SETUP_DONE, {}) setup_done_futures.update({domain: hass.loop.create_future() for domain in domains}) @@ -148,12 +148,8 @@ async def async_setup_component( if domain in hass.config.components: return True - setup_futures: dict[str, asyncio.Future[bool]] = hass.data.setdefault( - DATA_SETUP, {} - ) - setup_done_futures: dict[str, asyncio.Future[bool]] = hass.data.setdefault( - DATA_SETUP_DONE, {} - ) + setup_futures = hass.data.setdefault(DATA_SETUP, {}) + setup_done_futures = hass.data.setdefault(DATA_SETUP_DONE, {}) if existing_setup_future := setup_futures.get(domain): return await existing_setup_future @@ -194,9 +190,7 @@ async def _async_process_dependencies( Returns a list of dependencies which failed to set up. """ - setup_futures: dict[str, asyncio.Future[bool]] = hass.data.setdefault( - DATA_SETUP, {} - ) + setup_futures = hass.data.setdefault(DATA_SETUP, {}) dependencies_tasks = { dep: setup_futures.get(dep) @@ -209,7 +203,7 @@ async def _async_process_dependencies( } after_dependencies_tasks: dict[str, asyncio.Future[bool]] = {} - to_be_loaded: dict[str, asyncio.Future[bool]] = hass.data.get(DATA_SETUP_DONE, {}) + to_be_loaded = hass.data.get(DATA_SETUP_DONE, {}) for dep in integration.after_dependencies: if ( dep not in dependencies_tasks @@ -666,7 +660,7 @@ def _setup_started( """Return the setup started dict.""" if DATA_SETUP_STARTED not in hass.data: hass.data[DATA_SETUP_STARTED] = {} - return hass.data[DATA_SETUP_STARTED] # type: ignore[no-any-return] + return hass.data[DATA_SETUP_STARTED] @contextlib.contextmanager @@ -711,10 +705,10 @@ def _setup_times( ) -> defaultdict[str, defaultdict[str | None, defaultdict[SetupPhases, float]]]: """Return the setup timings default dict.""" if DATA_SETUP_TIME not in hass.data: - hass.data[DATA_SETUP_TIME] = defaultdict( + hass.data[DATA_SETUP_TIME] = defaultdict( # type: ignore[misc] lambda: defaultdict(lambda: defaultdict(float)) ) - return hass.data[DATA_SETUP_TIME] # type: ignore[no-any-return] + return hass.data[DATA_SETUP_TIME] @contextlib.contextmanager diff --git a/tests/test_setup.py b/tests/test_setup.py index e3d9a322862b6..5f2451a69886d 100644 --- a/tests/test_setup.py +++ b/tests/test_setup.py @@ -738,7 +738,6 @@ async def test_integration_only_setup_entry(hass: HomeAssistant) -> None: async def test_async_start_setup_running(hass: HomeAssistant) -> None: """Test setup started context manager does nothing when running.""" assert hass.state is CoreState.running - setup_started: dict[tuple[str, str | None], float] setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {}) with setup.async_start_setup( @@ -752,7 +751,6 @@ async def test_async_start_setup_config_entry( ) -> None: """Test setup started keeps track of setup times with a config entry.""" hass.set_state(CoreState.not_running) - setup_started: dict[tuple[str, str | None], float] setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {}) setup_time = setup._setup_times(hass) @@ -863,7 +861,6 @@ async def test_async_start_setup_config_entry_late_platform( ) -> None: """Test setup started tracks config entry time with a late platform load.""" hass.set_state(CoreState.not_running) - setup_started: dict[tuple[str, str | None], float] setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {}) setup_time = setup._setup_times(hass) @@ -918,7 +915,6 @@ async def test_async_start_setup_config_entry_platform_wait( ) -> None: """Test setup started tracks wait time when a platform loads inside of config entry setup.""" hass.set_state(CoreState.not_running) - setup_started: dict[tuple[str, str | None], float] setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {}) setup_time = setup._setup_times(hass) @@ -961,7 +957,6 @@ async def test_async_start_setup_config_entry_platform_wait( async def test_async_start_setup_top_level_yaml(hass: HomeAssistant) -> None: """Test setup started context manager keeps track of setup times with modern yaml.""" hass.set_state(CoreState.not_running) - setup_started: dict[tuple[str, str | None], float] setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {}) setup_time = setup._setup_times(hass) @@ -978,7 +973,6 @@ async def test_async_start_setup_top_level_yaml(hass: HomeAssistant) -> None: async def test_async_start_setup_platform_integration(hass: HomeAssistant) -> None: """Test setup started keeps track of setup times a platform integration.""" hass.set_state(CoreState.not_running) - setup_started: dict[tuple[str, str | None], float] setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {}) setup_time = setup._setup_times(hass) @@ -1013,7 +1007,6 @@ async def test_async_start_setup_legacy_platform_integration( ) -> None: """Test setup started keeps track of setup times for a legacy platform integration.""" hass.set_state(CoreState.not_running) - setup_started: dict[tuple[str, str | None], float] setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {}) setup_time = setup._setup_times(hass)