Use HassKey for hassio component data (#136172)

homeassistant/components/hassio/__init__.py

@@ -90,6 +90,7 @@ from .const import (
     ATTR_LOCATION,
     ATTR_PASSWORD,
     ATTR_SLUG,
+    DATA_COMPONENT,
     DATA_CORE_INFO,
     DATA_HOST_INFO,
     DATA_INFO,
@@ -326,7 +327,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:  # noqa:
 
     host = os.environ["SUPERVISOR"]
     websession = async_get_clientsession(hass)
-    hass.data[DOMAIN] = hassio = HassIO(hass.loop, websession, host)
+    hass.data[DATA_COMPONENT] = hassio = HassIO(hass.loop, websession, host)
     supervisor_client = get_supervisor_client(hass)
 
     try:

homeassistant/components/hassio/const.py

@@ -1,7 +1,16 @@
 """Hass.io const variables."""
 
+from __future__ import annotations
+
 from datetime import timedelta
 from enum import StrEnum
+from typing import TYPE_CHECKING
+
+from homeassistant.util.hass_dict import HassKey
+
+if TYPE_CHECKING:
+    from .handler import HassIO
+
 
 DOMAIN = "hassio"
 
@@ -64,6 +73,7 @@ UPDATE_KEY_SUPERVISOR = "supervisor"
 ADDONS_COORDINATOR = "hassio_addons_coordinator"
 
 
+DATA_COMPONENT: HassKey[HassIO] = HassKey(DOMAIN)
 DATA_CORE_INFO = "hassio_core_info"
 DATA_CORE_STATS = "hassio_core_stats"
 DATA_HOST_INFO = "hassio_host_info"

homeassistant/components/hassio/coordinator.py

@@ -35,6 +35,7 @@ from .const import (
     DATA_ADDONS_CHANGELOGS,
     DATA_ADDONS_INFO,
     DATA_ADDONS_STATS,
+    DATA_COMPONENT,
     DATA_CORE_INFO,
     DATA_CORE_STATS,
     DATA_HOST_INFO,
@@ -56,7 +57,7 @@ from .const import (
     SUPERVISOR_CONTAINER,
     SupervisorEntityModel,
 )
-from .handler import HassIO, HassioAPIError, get_supervisor_client
+from .handler import HassioAPIError, get_supervisor_client
 
 if TYPE_CHECKING:
     from .issues import SupervisorIssues
@@ -310,7 +311,7 @@ class HassioDataUpdateCoordinator(DataUpdateCoordinator):
                 hass, _LOGGER, cooldown=REQUEST_REFRESH_DELAY, immediate=False
             ),
         )
-        self.hassio: HassIO = hass.data[DOMAIN]
+        self.hassio = hass.data[DATA_COMPONENT]
         self.data = {}
         self.entry_id = config_entry.entry_id
         self.dev_reg = dev_reg

homeassistant/components/hassio/handler.py

@@ -24,7 +24,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.singleton import singleton
 from homeassistant.loader import bind_hass
 
-from .const import ATTR_MESSAGE, ATTR_RESULT, DOMAIN, X_HASS_SOURCE
+from .const import ATTR_MESSAGE, ATTR_RESULT, DATA_COMPONENT, X_HASS_SOURCE
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -72,7 +72,7 @@ async def async_update_diagnostics(hass: HomeAssistant, diagnostics: bool) -> bo
 
     The caller of the function should handle HassioAPIError.
     """
-    hassio: HassIO = hass.data[DOMAIN]
+    hassio = hass.data[DATA_COMPONENT]
     return await hassio.update_diagnostics(diagnostics)
 
 
@@ -85,7 +85,7 @@ async def async_create_backup(
 
     The caller of the function should handle HassioAPIError.
     """
-    hassio: HassIO = hass.data[DOMAIN]
+    hassio = hass.data[DATA_COMPONENT]
     backup_type = "partial" if partial else "full"
     command = f"/backups/new/{backup_type}"
     return await hassio.send_command(command, payload=payload, timeout=None)
@@ -94,7 +94,7 @@ async def async_create_backup(
 @api_data
 async def async_get_green_settings(hass: HomeAssistant) -> dict[str, bool]:
     """Return settings specific to Home Assistant Green."""
-    hassio: HassIO = hass.data[DOMAIN]
+    hassio = hass.data[DATA_COMPONENT]
     return await hassio.send_command("/os/boards/green", method="get")
 
 
@@ -106,7 +106,7 @@ async def async_set_green_settings(
 
     Returns an empty dict.
     """
-    hassio: HassIO = hass.data[DOMAIN]
+    hassio = hass.data[DATA_COMPONENT]
     return await hassio.send_command(
         "/os/boards/green", method="post", payload=settings
     )
@@ -115,7 +115,7 @@ async def async_set_green_settings(
 @api_data
 async def async_get_yellow_settings(hass: HomeAssistant) -> dict[str, bool]:
     """Return settings specific to Home Assistant Yellow."""
-    hassio: HassIO = hass.data[DOMAIN]
+    hassio = hass.data[DATA_COMPONENT]
     return await hassio.send_command("/os/boards/yellow", method="get")
 
 
@@ -127,7 +127,7 @@ async def async_set_yellow_settings(
 
     Returns an empty dict.
     """
-    hassio: HassIO = hass.data[DOMAIN]
+    hassio = hass.data[DATA_COMPONENT]
     return await hassio.send_command(
         "/os/boards/yellow", method="post", payload=settings
     )
@@ -333,7 +333,7 @@ class HassIO:
 @singleton(KEY_SUPERVISOR_CLIENT)
 def get_supervisor_client(hass: HomeAssistant) -> SupervisorClient:
     """Return supervisor client."""
-    hassio: HassIO = hass.data[DOMAIN]
+    hassio = hass.data[DATA_COMPONENT]
     return SupervisorClient(
         str(hassio.base_url),
         os.environ.get("SUPERVISOR_TOKEN", ""),

homeassistant/components/hassio/websocket_api.py

@@ -25,7 +25,7 @@ from .const import (
     ATTR_SESSION_DATA_USER_ID,
     ATTR_TIMEOUT,
     ATTR_WS_EVENT,
-    DOMAIN,
+    DATA_COMPONENT,
     EVENT_SUPERVISOR_EVENT,
     WS_ID,
     WS_TYPE,
@@ -33,7 +33,6 @@ from .const import (
     WS_TYPE_EVENT,
     WS_TYPE_SUBSCRIBE,
 )
-from .handler import HassIO
 
 SCHEMA_WEBSOCKET_EVENT = vol.Schema(
     {vol.Required(ATTR_WS_EVENT): cv.string},
@@ -113,7 +112,7 @@ async def websocket_supervisor_api(
         msg[ATTR_ENDPOINT]
     ):
         raise Unauthorized
-    supervisor: HassIO = hass.data[DOMAIN]
+    supervisor = hass.data[DATA_COMPONENT]
 
     command = msg[ATTR_ENDPOINT]
     payload = msg.get(ATTR_DATA, {})