diff --git a/homeassistant/components/trace/__init__.py b/homeassistant/components/trace/__init__.py index bb4f046a7e2..3d7510b57b2 100644 --- a/homeassistant/components/trace/__init__.py +++ b/homeassistant/components/trace/__init__.py @@ -1,12 +1,14 @@ """Support for script and automation tracing and debugging.""" from __future__ import annotations +from collections.abc import Mapping import logging +from typing import Any import voluptuous as vol from homeassistant.const import EVENT_HOMEASSISTANT_STOP -from homeassistant.core import HomeAssistant +from homeassistant.core import Event, HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError import homeassistant.helpers.config_validation as cv from homeassistant.helpers.json import ExtendedJSONEncoder @@ -21,7 +23,7 @@ from .const import ( DATA_TRACES_RESTORED, DEFAULT_STORED_TRACES, ) -from .models import ActionTrace, BaseTrace, RestoredTrace # noqa: F401 +from .models import ActionTrace, BaseTrace, RestoredTrace from .utils import LimitedSizeDict _LOGGER = logging.getLogger(__name__) @@ -35,6 +37,13 @@ TRACE_CONFIG_SCHEMA = { vol.Optional(CONF_STORED_TRACES, default=DEFAULT_STORED_TRACES): cv.positive_int } +TraceData = dict[str, LimitedSizeDict[str, BaseTrace]] + + +@callback +def _get_data(hass: HomeAssistant) -> TraceData: + return hass.data[DATA_TRACE] + async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Initialize the trace integration.""" @@ -45,15 +54,12 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: ) hass.data[DATA_TRACE_STORE] = store - async def _async_store_traces_at_stop(*_) -> None: + async def _async_store_traces_at_stop(_: Event) -> None: """Save traces to storage.""" _LOGGER.debug("Storing traces") try: await store.async_save( - { - key: list(traces.values()) - for key, traces in hass.data[DATA_TRACE].items() - } + {key: list(traces.values()) for key, traces in _get_data(hass).items()} ) except HomeAssistantError as exc: _LOGGER.error("Error storing traces", exc_info=exc) @@ -64,25 +70,30 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: return True -async def async_get_trace(hass, key, run_id): +async def async_get_trace( + hass: HomeAssistant, key: str, run_id: str +) -> dict[str, BaseTrace]: """Return the requested trace.""" # Restore saved traces if not done await async_restore_traces(hass) - return hass.data[DATA_TRACE][key][run_id].as_extended_dict() + return _get_data(hass)[key][run_id].as_extended_dict() -async def async_list_contexts(hass, key): +async def async_list_contexts( + hass: HomeAssistant, key: str | None +) -> dict[str, dict[str, str]]: """List contexts for which we have traces.""" # Restore saved traces if not done await async_restore_traces(hass) + values: Mapping[str, LimitedSizeDict[str, BaseTrace] | None] if key is not None: - values = {key: hass.data[DATA_TRACE].get(key, {})} + values = {key: _get_data(hass).get(key)} else: - values = hass.data[DATA_TRACE] + values = _get_data(hass) - def _trace_id(run_id, key) -> dict: + def _trace_id(run_id: str, key: str) -> dict[str, str]: """Make trace_id for the response.""" domain, item_id = key.split(".", 1) return {"run_id": run_id, "domain": domain, "item_id": item_id} @@ -90,28 +101,32 @@ async def async_list_contexts(hass, key): return { trace.context.id: _trace_id(trace.run_id, key) for key, traces in values.items() + if traces is not None for trace in traces.values() } -def _get_debug_traces(hass, key): +def _get_debug_traces(hass: HomeAssistant, key: str) -> list[dict[str, Any]]: """Return a serializable list of debug traces for a script or automation.""" - traces = [] + traces: list[dict[str, Any]] = [] - for trace in hass.data[DATA_TRACE].get(key, {}).values(): - traces.append(trace.as_short_dict()) + if traces_for_key := _get_data(hass).get(key): + for trace in traces_for_key.values(): + traces.append(trace.as_short_dict()) return traces -async def async_list_traces(hass, wanted_domain, wanted_key): +async def async_list_traces( + hass: HomeAssistant, wanted_domain: str, wanted_key: str | None +) -> list[dict[str, Any]]: """List traces for a domain.""" # Restore saved traces if not done already await async_restore_traces(hass) if not wanted_key: - traces = [] - for key in hass.data[DATA_TRACE]: + traces: list[dict[str, Any]] = [] + for key in _get_data(hass): domain = key.split(".", 1)[0] if domain == wanted_domain: traces.extend(_get_debug_traces(hass, key)) @@ -126,7 +141,7 @@ def async_store_trace( ) -> None: """Store a trace if its key is valid.""" if key := trace.key: - traces = hass.data[DATA_TRACE] + traces = _get_data(hass) if key not in traces: traces[key] = LimitedSizeDict(size_limit=stored_traces) else: @@ -137,7 +152,7 @@ def async_store_trace( def _async_store_restored_trace(hass: HomeAssistant, trace: RestoredTrace) -> None: """Store a restored trace and move it to the end of the LimitedSizeDict.""" key = trace.key - traces = hass.data[DATA_TRACE] + traces = _get_data(hass) if key not in traces: traces[key] = LimitedSizeDict() traces[key][trace.run_id] = trace @@ -151,7 +166,7 @@ async def async_restore_traces(hass: HomeAssistant) -> None: hass.data[DATA_TRACES_RESTORED] = True - store = hass.data[DATA_TRACE_STORE] + store: Store[dict[str, list]] = hass.data[DATA_TRACE_STORE] try: restored_traces = await store.async_load() or {} except HomeAssistantError: @@ -162,7 +177,7 @@ async def async_restore_traces(hass: HomeAssistant) -> None: # Add stored traces in reversed order to priorize the newest traces for json_trace in reversed(traces): if ( - (stored_traces := hass.data[DATA_TRACE].get(key)) + (stored_traces := _get_data(hass).get(key)) and stored_traces.size_limit is not None and len(stored_traces) >= stored_traces.size_limit ):