diff --git a/homeassistant/helpers/script.py b/homeassistant/helpers/script.py index be6be32da1ea..ac56ce2beb42 100644 --- a/homeassistant/helpers/script.py +++ b/homeassistant/helpers/script.py @@ -72,6 +72,7 @@ from homeassistant.util.dt import utcnow from .trace import ( TraceElement, + async_trace_path, trace_append_element, trace_id_get, trace_path, @@ -613,11 +614,14 @@ class _ScriptRun: if not check: raise _StopScript - def _test_conditions(self, conditions, name): + def _test_conditions(self, conditions, name, condition_path=None): + if condition_path is None: + condition_path = name + @trace_condition_function def traced_test_conditions(hass, variables): try: - with trace_path("conditions"): + with trace_path(condition_path): for idx, cond in enumerate(conditions): with trace_path(str(idx)): if not cond(hass, variables): @@ -631,7 +635,7 @@ class _ScriptRun: result = traced_test_conditions(self._hass, self._variables) return result - @trace_path("repeat") + @async_trace_path("repeat") async def _async_repeat_step(self): """Repeat a sequence.""" description = self._action.get(CONF_ALIAS, "sequence") @@ -720,7 +724,7 @@ class _ScriptRun: for idx, (conditions, script) in enumerate(choose_data["choices"]): with trace_path(str(idx)): try: - if self._test_conditions(conditions, "choose"): + if self._test_conditions(conditions, "choose", "conditions"): trace_set_result(choice=idx) with trace_path("sequence"): await self._async_run_script(script) diff --git a/homeassistant/helpers/trace.py b/homeassistant/helpers/trace.py index 2e5f1dd8c540..d6de845248f1 100644 --- a/homeassistant/helpers/trace.py +++ b/homeassistant/helpers/trace.py @@ -4,7 +4,8 @@ from __future__ import annotations from collections import deque from contextlib import contextmanager from contextvars import ContextVar -from typing import Any, Deque, Generator, cast +from functools import wraps +from typing import Any, Callable, Deque, Generator, cast from homeassistant.helpers.typing import TemplateVarsType import homeassistant.util.dt as dt_util @@ -168,9 +169,32 @@ def trace_set_result(**kwargs: Any) -> None: @contextmanager def trace_path(suffix: str | list[str]) -> Generator: - """Go deeper in the config tree.""" + """Go deeper in the config tree. + + Can not be used as a decorator on couroutine functions. + """ count = trace_path_push(suffix) try: yield finally: trace_path_pop(count) + + +def async_trace_path(suffix: str | list[str]) -> Callable: + """Go deeper in the config tree. + + To be used as a decorator on coroutine functions. + """ + + def _trace_path_decorator(func: Callable) -> Callable: + """Decorate a coroutine function.""" + + @wraps(func) + async def async_wrapper(*args: Any) -> None: + """Catch and log exception.""" + with trace_path(suffix): + await func(*args) + + return async_wrapper + + return _trace_path_decorator diff --git a/tests/helpers/test_script.py b/tests/helpers/test_script.py index f62462990749..1f47aa9dbbaf 100644 --- a/tests/helpers/test_script.py +++ b/tests/helpers/test_script.py @@ -1236,7 +1236,7 @@ async def test_repeat_count(hass, caplog, count): assert_action_trace( { "0": [{}], - "0/sequence/0": [{}] * min(count, script.ACTION_TRACE_NODE_MAX_LEN), + "0/repeat/sequence/0": [{}] * min(count, script.ACTION_TRACE_NODE_MAX_LEN), } )