Address review comments from trace refactoring PRs

homeassistant/components/automation/trace.py

@@ -2,25 +2,79 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
+from typing import Any, Deque
 
-from homeassistant.components.trace import AutomationTrace, async_store_trace
+from homeassistant.components.trace import ActionTrace, async_store_trace
+from homeassistant.core import Context
+from homeassistant.helpers.trace import TraceElement
 
 # mypy: allow-untyped-calls, allow-untyped-defs
 # mypy: no-check-untyped-defs, no-warn-return-any
 
 
+class AutomationTrace(ActionTrace):
+    """Container for automation trace."""
+
+    def __init__(
+        self,
+        item_id: str,
+        config: dict[str, Any],
+        context: Context,
+    ):
+        """Container for automation trace."""
+        key = ("automation", item_id)
+        super().__init__(key, config, context)
+        self._condition_trace: dict[str, Deque[TraceElement]] | None = None
+
+    def set_condition_trace(self, trace: dict[str, Deque[TraceElement]]) -> None:
+        """Set condition trace."""
+        self._condition_trace = trace
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return dictionary version of this AutomationTrace."""
+
+        result = super().as_dict()
+
+        condition_traces = {}
+
+        if self._condition_trace:
+            for key, trace_list in self._condition_trace.items():
+                condition_traces[key] = [item.as_dict() for item in trace_list]
+        result["condition_trace"] = condition_traces
+
+        return result
+
+    def as_short_dict(self) -> dict[str, Any]:
+        """Return a brief dictionary version of this AutomationTrace."""
+
+        result = super().as_short_dict()
+
+        last_condition = None
+        trigger = None
+
+        if self._condition_trace:
+            last_condition = list(self._condition_trace)[-1]
+        if self._variables:
+            trigger = self._variables.get("trigger", {}).get("description")
+
+        result["trigger"] = trigger
+        result["last_condition"] = last_condition
+
+        return result
+
+
 @contextmanager
-def trace_automation(hass, item_id, config, context):
-    """Trace action execution of automation with item_id."""
-    trace = AutomationTrace(item_id, config, context)
+def trace_automation(hass, automation_id, config, context):
+    """Trace action execution of automation with automation_id."""
+    trace = AutomationTrace(automation_id, config, context)
     async_store_trace(hass, trace)
 
     try:
         yield trace
     except Exception as ex:  # pylint: disable=broad-except
-        if item_id:
+        if automation_id:
             trace.set_error(ex)
         raise ex
     finally:
-        if item_id:
+        if automation_id:
             trace.finished()

homeassistant/components/script/trace.py

@@ -2,8 +2,24 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
+from typing import Any
 
-from homeassistant.components.trace import ScriptTrace, async_store_trace
+from homeassistant.components.trace import ActionTrace, async_store_trace
+from homeassistant.core import Context
+
+
+class ScriptTrace(ActionTrace):
+    """Container for automation trace."""
+
+    def __init__(
+        self,
+        item_id: str,
+        config: dict[str, Any],
+        context: Context,
+    ):
+        """Container for automation trace."""
+        key = ("script", item_id)
+        super().__init__(key, config, context)
 
 
 @contextmanager

homeassistant/components/trace/__init__.py

@@ -121,68 +121,3 @@ def as_short_dict(self) -> dict[str, Any]:
             result["last_action"] = last_action
 
         return result
-
-
-class AutomationTrace(ActionTrace):
-    """Container for automation trace."""
-
-    def __init__(
-        self,
-        item_id: str,
-        config: dict[str, Any],
-        context: Context,
-    ):
-        """Container for automation trace."""
-        key = ("automation", item_id)
-        super().__init__(key, config, context)
-        self._condition_trace: dict[str, Deque[TraceElement]] | None = None
-
-    def set_condition_trace(self, trace: dict[str, Deque[TraceElement]]) -> None:
-        """Set condition trace."""
-        self._condition_trace = trace
-
-    def as_dict(self) -> dict[str, Any]:
-        """Return dictionary version of this AutomationTrace."""
-
-        result = super().as_dict()
-
-        condition_traces = {}
-
-        if self._condition_trace:
-            for key, trace_list in self._condition_trace.items():
-                condition_traces[key] = [item.as_dict() for item in trace_list]
-        result["condition_trace"] = condition_traces
-
-        return result
-
-    def as_short_dict(self) -> dict[str, Any]:
-        """Return a brief dictionary version of this AutomationTrace."""
-
-        result = super().as_short_dict()
-
-        last_condition = None
-        trigger = None
-
-        if self._condition_trace:
-            last_condition = list(self._condition_trace)[-1]
-        if self._variables:
-            trigger = self._variables.get("trigger", {}).get("description")
-
-        result["trigger"] = trigger
-        result["last_condition"] = last_condition
-
-        return result
-
-
-class ScriptTrace(ActionTrace):
-    """Container for automation trace."""
-
-    def __init__(
-        self,
-        item_id: str,
-        config: dict[str, Any],
-        context: Context,
-    ):
-        """Container for automation trace."""
-        key = ("script", item_id)
-        super().__init__(key, config, context)

homeassistant/components/trace/utils.py

@@ -1,4 +1,4 @@
-"""Helpers for automation and script tracing and debugging."""
+"""Helpers for script and automation tracing and debugging."""
 from collections import OrderedDict
 from datetime import timedelta
 from typing import Any

homeassistant/components/trace/websocket_api.py

@@ -23,12 +23,12 @@
     debug_stop,
 )
 
-from .trace import DATA_TRACE, get_all_debug_traces, get_debug_trace, get_debug_traces
+from .const import DATA_TRACE
 from .utils import TraceJSONEncoder
 
 # mypy: allow-untyped-calls, allow-untyped-defs
 
-TRACE_DOMAINS = ["automation", "script"]
+TRACE_DOMAINS = ("automation", "script")
 
 
 @callback
@@ -57,33 +57,47 @@ def async_setup(hass: HomeAssistant) -> None:
     }
 )
 def websocket_trace_get(hass, connection, msg):
-    """Get an automation or script trace."""
+    """Get an script or automation trace."""
     key = (msg["domain"], msg["item_id"])
     run_id = msg["run_id"]
 
-    trace = get_debug_trace(hass, key, run_id)
+    trace = hass.data[DATA_TRACE][key][run_id]
     message = websocket_api.messages.result_message(msg["id"], trace)
 
     connection.send_message(json.dumps(message, cls=TraceJSONEncoder, allow_nan=False))
 
 
+def get_debug_traces(hass, key):
+    """Return a serializable list of debug traces for an script or automation."""
+    traces = []
+
+    for trace in hass.data[DATA_TRACE].get(key, {}).values():
+        traces.append(trace.as_short_dict())
+
+    return traces
+
+
 @callback
 @websocket_api.require_admin
 @websocket_api.websocket_command(
     {
         vol.Required("type"): "trace/list",
-        vol.Inclusive("domain", "id"): vol.In(TRACE_DOMAINS),
-        vol.Inclusive("item_id", "id"): str,
+        vol.Required("domain", "id"): vol.In(TRACE_DOMAINS),
+        vol.Optional("item_id", "id"): str,
     }
 )
 def websocket_trace_list(hass, connection, msg):
-    """Summarize automation and script traces."""
-    key = (msg["domain"], msg["item_id"]) if "item_id" in msg else None
+    """Summarize script and automation traces."""
+    domain = msg["domain"]
+    key = (domain, msg["item_id"]) if "item_id" in msg else None
 
     if not key:
-        traces = get_all_debug_traces(hass, summary=True)
+        traces = []
+        for key in hass.data[DATA_TRACE]:
+            if key[0] == domain:
+                traces.extend(get_debug_traces(hass, key))
     else:
-        traces = get_debug_traces(hass, key, summary=True)
+        traces = get_debug_traces(hass, key)
 
     connection.send_result(msg["id"], traces)
 
@@ -230,7 +244,7 @@ def unsub():
     }
 )
 def websocket_debug_continue(hass, connection, msg):
-    """Resume execution of halted automation or script."""
+    """Resume execution of halted script or automation."""
     key = (msg["domain"], msg["item_id"])
     run_id = msg["run_id"]
 
@@ -250,7 +264,7 @@ def websocket_debug_continue(hass, connection, msg):
     }
 )
 def websocket_debug_step(hass, connection, msg):
-    """Single step a halted automation or script."""
+    """Single step a halted script or automation."""
     key = (msg["domain"], msg["item_id"])
     run_id = msg["run_id"]
 
@@ -270,7 +284,7 @@ def websocket_debug_step(hass, connection, msg):
     }
 )
 def websocket_debug_stop(hass, connection, msg):
-    """Stop a halted automation or script."""
+    """Stop a halted script or automation."""
     key = (msg["domain"], msg["item_id"])
     run_id = msg["run_id"]