Set variable values in scripts (#39915)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/helpers/config_validation.py

@@ -67,6 +67,7 @@ from homeassistant.const import (
     CONF_UNIT_SYSTEM_METRIC,
     CONF_UNTIL,
     CONF_VALUE_TEMPLATE,
+    CONF_VARIABLES,
     CONF_WAIT_FOR_TRIGGER,
     CONF_WAIT_TEMPLATE,
     CONF_WHILE,
@@ -1127,6 +1128,13 @@ _SCRIPT_WAIT_FOR_TRIGGER_SCHEMA = vol.Schema(
     }
 )
 
+_SCRIPT_SET_SCHEMA = vol.Schema(
+    {
+        vol.Optional(CONF_ALIAS): string,
+        vol.Required(CONF_VARIABLES): SCRIPT_VARIABLES_SCHEMA,
+    }
+)
+
 SCRIPT_ACTION_DELAY = "delay"
 SCRIPT_ACTION_WAIT_TEMPLATE = "wait_template"
 SCRIPT_ACTION_CHECK_CONDITION = "condition"
@@ -1137,6 +1145,7 @@ SCRIPT_ACTION_ACTIVATE_SCENE = "scene"
 SCRIPT_ACTION_REPEAT = "repeat"
 SCRIPT_ACTION_CHOOSE = "choose"
 SCRIPT_ACTION_WAIT_FOR_TRIGGER = "wait_for_trigger"
+SCRIPT_ACTION_VARIABLES = "variables"
 
 
 def determine_script_action(action: dict) -> str:
@@ -1168,6 +1177,9 @@ def determine_script_action(action: dict) -> str:
     if CONF_WAIT_FOR_TRIGGER in action:
         return SCRIPT_ACTION_WAIT_FOR_TRIGGER
 
+    if CONF_VARIABLES in action:
+        return SCRIPT_ACTION_VARIABLES
+
     return SCRIPT_ACTION_CALL_SERVICE
 
 
@@ -1182,4 +1194,5 @@ ACTION_TYPE_SCHEMAS: Dict[str, Callable[[Any], dict]] = {
     SCRIPT_ACTION_REPEAT: _SCRIPT_REPEAT_SCHEMA,
     SCRIPT_ACTION_CHOOSE: _SCRIPT_CHOOSE_SCHEMA,
     SCRIPT_ACTION_WAIT_FOR_TRIGGER: _SCRIPT_WAIT_FOR_TRIGGER_SCHEMA,
+    SCRIPT_ACTION_VARIABLES: _SCRIPT_SET_SCHEMA,
 }

homeassistant/helpers/script.py

@@ -46,6 +46,7 @@ from homeassistant.const import (
     CONF_SEQUENCE,
     CONF_TIMEOUT,
     CONF_UNTIL,
+    CONF_VARIABLES,
     CONF_WAIT_FOR_TRIGGER,
     CONF_WAIT_TEMPLATE,
     CONF_WHILE,
@@ -612,6 +613,14 @@ class _ScriptRun:
                 task.cancel()
             remove_triggers()
 
+    async def _async_variables_step(self):
+        """Set a variable value."""
+        self._script.last_action = self._action.get(CONF_ALIAS, "setting variables")
+        self._log("Executing step %s", self._script.last_action)
+        self._variables = self._action[CONF_VARIABLES].async_render(
+            self._hass, self._variables, render_as_defaults=False
+        )
+
     async def _async_run_script(self, script):
         """Execute a script."""
         await self._async_run_long_action(

homeassistant/helpers/script_variables.py

@@ -19,21 +19,31 @@ class ScriptVariables:
         self,
         hass: HomeAssistant,
         run_variables: Optional[Mapping[str, Any]],
+        *,
+        render_as_defaults: bool = True,
     ) -> Dict[str, Any]:
         """Render script variables.
 
-        The run variables are used to compute the static variables, but afterwards will also
-        be merged on top of the static variables.
+        The run variables are used to compute the static variables.
+
+        If `render_as_defaults` is True, the run variables will not be overridden.
+
         """
         if self._has_template is None:
             self._has_template = template.is_complex(self.variables)
             template.attach(hass, self.variables)
 
         if not self._has_template:
-            rendered_variables = dict(self.variables)
+            if render_as_defaults:
+                rendered_variables = dict(self.variables)
 
-            if run_variables is not None:
-                rendered_variables.update(run_variables)
+                if run_variables is not None:
+                    rendered_variables.update(run_variables)
+            else:
+                rendered_variables = (
+                    {} if run_variables is None else dict(run_variables)
+                )
+                rendered_variables.update(self.variables)
 
             return rendered_variables
 
@@ -42,14 +52,11 @@ class ScriptVariables:
         for key, value in self.variables.items():
             # We can skip if we're going to override this key with
             # run variables anyway
-            if key in rendered_variables:
+            if render_as_defaults and key in rendered_variables:
                 continue
 
             rendered_variables[key] = template.render_complex(value, rendered_variables)
 
-        if run_variables:
-            rendered_variables.update(run_variables)
-
         return rendered_variables
 
     def as_dict(self) -> dict:

tests/helpers/test_script.py

@@ -1785,3 +1785,42 @@ async def test_started_action(hass, caplog):
     await hass.async_block_till_done()
 
     assert log_message in caplog.text
+
+
+async def test_set_variable(hass, caplog):
+    """Test setting variables in scripts."""
+    sequence = cv.SCRIPT_SCHEMA(
+        [
+            {"variables": {"variable": "value"}},
+            {"service": "test.script", "data": {"value": "{{ variable }}"}},
+        ]
+    )
+    script_obj = script.Script(hass, sequence, "test script", "test_domain")
+
+    mock_calls = async_mock_service(hass, "test", "script")
+
+    await script_obj.async_run(context=Context())
+    await hass.async_block_till_done()
+
+    assert mock_calls[0].data["value"] == "value"
+
+
+async def test_set_redefines_variable(hass, caplog):
+    """Test setting variables based on their current value."""
+    sequence = cv.SCRIPT_SCHEMA(
+        [
+            {"variables": {"variable": "1"}},
+            {"service": "test.script", "data": {"value": "{{ variable }}"}},
+            {"variables": {"variable": "{{ variable | int + 1 }}"}},
+            {"service": "test.script", "data": {"value": "{{ variable }}"}},
+        ]
+    )
+    script_obj = script.Script(hass, sequence, "test script", "test_domain")
+
+    mock_calls = async_mock_service(hass, "test", "script")
+
+    await script_obj.async_run(context=Context())
+    await hass.async_block_till_done()
+
+    assert mock_calls[0].data["value"] == "1"
+    assert mock_calls[1].data["value"] == "2"

tests/helpers/test_script_variables.py

@@ -24,6 +24,28 @@ async def test_static_vars_run_args():
     assert orig == orig_copy
 
 
+async def test_static_vars_no_default():
+    """Test static vars."""
+    orig = {"hello": "world"}
+    var = cv.SCRIPT_VARIABLES_SCHEMA(orig)
+    rendered = var.async_render(None, None, render_as_defaults=False)
+    assert rendered is not orig
+    assert rendered == orig
+
+
+async def test_static_vars_run_args_no_default():
+    """Test static vars."""
+    orig = {"hello": "world"}
+    orig_copy = dict(orig)
+    var = cv.SCRIPT_VARIABLES_SCHEMA(orig)
+    rendered = var.async_render(
+        None, {"hello": "override", "run": "var"}, render_as_defaults=False
+    )
+    assert rendered == {"hello": "world", "run": "var"}
+    # Make sure we don't change original vars
+    assert orig == orig_copy
+
+
 async def test_template_vars(hass):
     """Test template vars."""
     var = cv.SCRIPT_VARIABLES_SCHEMA({"hello": "{{ 1 + 1 }}"})
@@ -53,6 +75,36 @@ async def test_template_vars_run_args(hass):
     }
 
 
+async def test_template_vars_no_default(hass):
+    """Test template vars."""
+    var = cv.SCRIPT_VARIABLES_SCHEMA({"hello": "{{ 1 + 1 }}"})
+    rendered = var.async_render(hass, None, render_as_defaults=False)
+    assert rendered == {"hello": "2"}
+
+
+async def test_template_vars_run_args_no_default(hass):
+    """Test template vars."""
+    var = cv.SCRIPT_VARIABLES_SCHEMA(
+        {
+            "something": "{{ run_var_ex + 1 }}",
+            "something_2": "{{ run_var_ex + 1 }}",
+        }
+    )
+    rendered = var.async_render(
+        hass,
+        {
+            "run_var_ex": 5,
+            "something_2": 1,
+        },
+        render_as_defaults=False,
+    )
+    assert rendered == {
+        "run_var_ex": 5,
+        "something": "6",
+        "something_2": "6",
+    }
+
+
 async def test_template_vars_error(hass):
     """Test template vars."""
     var = cv.SCRIPT_VARIABLES_SCHEMA({"hello": "{{ canont.work }}"})