Use HassKey for setup keys

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()
 

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.

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

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)