From 6278d36981285ffe932463f606ff7ca9d7987432 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Tue, 28 Jan 2025 13:49:49 +0100 Subject: [PATCH] Use HassKey in diagnostics (#136627) --- homeassistant/components/diagnostics/__init__.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/homeassistant/components/diagnostics/__init__.py b/homeassistant/components/diagnostics/__init__.py index b23b7cef2bd..7bc43f2c3f5 100644 --- a/homeassistant/components/diagnostics/__init__.py +++ b/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)