Enable augmented-assignment operations in scripts (#108081)

homeassistant/components/python_script/__init__.py

@@ -2,8 +2,11 @@
 import datetime
 import glob
 import logging
+from numbers import Number
+import operator
 import os
 import time
+from typing import Any
 
 from RestrictedPython import (
     compile_restricted_exec,
@@ -146,6 +149,36 @@ def discover_scripts(hass):
         async_set_service_schema(hass, DOMAIN, name, service_desc)
 
 
+IOPERATOR_TO_OPERATOR = {
+    "%=": operator.mod,
+    "&=": operator.and_,
+    "**=": operator.pow,
+    "*=": operator.mul,
+    "+=": operator.add,
+    "-=": operator.sub,
+    "//=": operator.floordiv,
+    "/=": operator.truediv,
+    "<<=": operator.lshift,
+    ">>=": operator.rshift,
+    "@=": operator.matmul,
+    "^=": operator.xor,
+    "|=": operator.or_,
+}
+
+
+def guarded_inplacevar(op: str, target: Any, operand: Any) -> Any:
+    """Implement augmented-assign (+=, -=, etc.) operators for restricted code.
+
+    See RestrictedPython's `visit_AugAssign` for details.
+    """
+    if not isinstance(target, (list, Number, str)):
+        raise ScriptError(f"The {op!r} operation is not allowed on a {type(target)}")
+    op_fun = IOPERATOR_TO_OPERATOR.get(op)
+    if not op_fun:
+        raise ScriptError(f"The {op!r} operation is not allowed")
+    return op_fun(target, operand)
+
+
 @bind_hass
 def execute_script(hass, name, data=None, return_response=False):
     """Execute a script."""
@@ -223,6 +256,7 @@ def execute(hass, filename, source, data=None, return_response=False):
         "_getitem_": default_guarded_getitem,
         "_iter_unpack_sequence_": guarded_iter_unpack_sequence,
         "_unpack_sequence_": guarded_unpack_sequence,
+        "_inplacevar_": guarded_inplacevar,
         "hass": hass,
         "data": data or {},
         "logger": logger,