Use HassKey in diagnostics (#136627)

homeassistant/components/diagnostics/__init__.py

@@ -33,6 +33,7 @@ from homeassistant.loader import (
     async_get_integration,
 )
 from homeassistant.setup import async_get_domain_setup_times
+from homeassistant.util.hass_dict import HassKey
 from homeassistant.util.json import format_unserializable_data
 
 from .const import DOMAIN, REDACTED, DiagnosticsSubType, DiagnosticsType
@@ -44,6 +45,7 @@ _LOGGER = logging.getLogger(__name__)
 
 
 CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN)
+_DIAGNOSTICS_DATA: HassKey[DiagnosticsData] = HassKey(DOMAIN)
 
 
 @dataclass(slots=True)
@@ -72,7 +74,7 @@ class DiagnosticsData:
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up Diagnostics from a config entry."""
-    hass.data[DOMAIN] = DiagnosticsData()
+    hass.data[_DIAGNOSTICS_DATA] = DiagnosticsData()
 
     await integration_platform.async_process_integration_platforms(
         hass, DOMAIN, _register_diagnostics_platform
@@ -104,7 +106,7 @@ def _register_diagnostics_platform(
     hass: HomeAssistant, integration_domain: str, platform: DiagnosticsProtocol
 ) -> None:
     """Register a diagnostics platform."""
-    diagnostics_data: DiagnosticsData = hass.data[DOMAIN]
+    diagnostics_data = hass.data[_DIAGNOSTICS_DATA]
     diagnostics_data.platforms[integration_domain] = DiagnosticsPlatformData(
         getattr(platform, "async_get_config_entry_diagnostics", None),
         getattr(platform, "async_get_device_diagnostics", None),
@@ -118,7 +120,7 @@ def handle_info(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict
 ) -> None:
     """List all possible diagnostic handlers."""
-    diagnostics_data: DiagnosticsData = hass.data[DOMAIN]
+    diagnostics_data = hass.data[_DIAGNOSTICS_DATA]
     result = [
         {
             "domain": domain,
@@ -145,7 +147,7 @@ def handle_get(
 ) -> None:
     """List all diagnostic handlers for a domain."""
     domain = msg["domain"]
-    diagnostics_data: DiagnosticsData = hass.data[DOMAIN]
+    diagnostics_data = hass.data[_DIAGNOSTICS_DATA]
 
     if (info := diagnostics_data.platforms.get(domain)) is None:
         connection.send_error(
@@ -267,7 +269,7 @@ class DownloadDiagnosticsView(http.HomeAssistantView):
         if (config_entry := hass.config_entries.async_get_entry(d_id)) is None:
             return web.Response(status=HTTPStatus.NOT_FOUND)
 
-        diagnostics_data: DiagnosticsData = hass.data[DOMAIN]
+        diagnostics_data = hass.data[_DIAGNOSTICS_DATA]
         if (info := diagnostics_data.platforms.get(config_entry.domain)) is None:
             return web.Response(status=HTTPStatus.NOT_FOUND)