Refactor tracing: Prepare for tracing of scripts

homeassistant/components/automation/trace.py

@@ -23,7 +23,7 @@ class AutomationTrace:
 
     def __init__(
         self,
-        unique_id: str | None,
+        key: tuple[str, str],
         config: dict[str, Any],
         context: Context,
     ):
@@ -37,7 +37,7 @@ def __init__(
         self.run_id: str = str(next(self._run_ids))
         self._timestamp_finish: dt.datetime | None = None
         self._timestamp_start: dt.datetime = dt_util.utcnow()
-        self._unique_id: str | None = unique_id
+        self._key: tuple[str, str] = key
         self._variables: dict[str, Any] | None = None
 
     def set_action_trace(self, trace: dict[str, Deque[TraceElement]]) -> None:
@@ -104,7 +104,6 @@ def as_short_dict(self) -> dict[str, Any]:
             trigger = self._variables.get("trigger", {}).get("description")
 
         result = {
-            "automation_id": self._unique_id,
             "last_action": last_action,
             "last_condition": last_condition,
             "run_id": self.run_id,
@@ -114,7 +113,8 @@ def as_short_dict(self) -> dict[str, Any]:
                 "finish": self._timestamp_finish,
             },
             "trigger": trigger,
-            "unique_id": self._unique_id,
+            "domain": self._key[0],
+            "item_id": self._key[1],
         }
         if self._error is not None:
             result["error"] = str(self._error)
@@ -126,23 +126,24 @@ def as_short_dict(self) -> dict[str, Any]:
 
 
 @contextmanager
-def trace_automation(hass, unique_id, config, context):
+def trace_automation(hass, item_id, config, context):
     """Trace action execution of automation with automation_id."""
-    automation_trace = AutomationTrace(unique_id, config, context)
-    trace_id_set((unique_id, automation_trace.run_id))
+    key = ("automation", item_id)
+    trace = AutomationTrace(key, config, context)
+    trace_id_set((key, trace.run_id))
 
-    if unique_id:
-        automation_traces = hass.data[DATA_TRACE]
-        if unique_id not in automation_traces:
-            automation_traces[unique_id] = LimitedSizeDict(size_limit=STORED_TRACES)
-        automation_traces[unique_id][automation_trace.run_id] = automation_trace
+    if key:
+        traces = hass.data[DATA_TRACE]
+        if key not in traces:
+            traces[key] = LimitedSizeDict(size_limit=STORED_TRACES)
+        traces[key][trace.run_id] = trace
 
     try:
-        yield automation_trace
+        yield trace
     except Exception as ex:  # pylint: disable=broad-except
-        if unique_id:
-            automation_trace.set_error(ex)
+        if key:
+            trace.set_error(ex)
         raise ex
     finally:
-        if unique_id:
-            automation_trace.finished()
+        if key:
+            trace.finished()

homeassistant/components/trace/trace.py

@@ -5,17 +5,17 @@
 
 
 @callback
-def get_debug_trace(hass, automation_id, run_id):
+def get_debug_trace(hass, key, run_id):
     """Return a serializable debug trace."""
-    return hass.data[DATA_TRACE][automation_id][run_id]
+    return hass.data[DATA_TRACE][key][run_id]
 
 
 @callback
-def get_debug_traces_for_automation(hass, automation_id, summary=False):
-    """Return a serializable list of debug traces for an automation."""
+def get_debug_traces(hass, key, summary=False):
+    """Return a serializable list of debug traces for an automation or script."""
     traces = []
 
-    for trace in hass.data[DATA_TRACE].get(automation_id, {}).values():
+    for trace in hass.data[DATA_TRACE].get(key, {}).values():
         if summary:
             traces.append(trace.as_short_dict())
         else:
@@ -25,11 +25,11 @@ def get_debug_traces_for_automation(hass, automation_id, summary=False):
 
 
 @callback
-def get_debug_traces(hass, summary=False):
-    """Return a serializable list of debug traces."""
+def get_all_debug_traces(hass, summary=False):
+    """Return a serializable list of debug traces for all automations and scripts."""
     traces = []
 
-    for automation_id in hass.data[DATA_TRACE]:
-        traces.extend(get_debug_traces_for_automation(hass, automation_id, summary))
+    for key in hass.data[DATA_TRACE]:
+        traces.extend(get_debug_traces(hass, key, summary))
 
     return traces