Support assignment expressions (i.e. the ":=" operator) (#195)

* Support assignment expressions (i.e. the ":=" operator

* make flake8 happy

* add tests for `NamedExpr` support

* make `isort` happy

* check that "private" variables cannot be used as target

* use version detection support from `_compat`;
get rid of `util`

* make `isort` happy

* move `test_NamedExpr` out of `src/RestrictedPython/tests` to `tests`

Co-authored-by: Jens Vagelpohl <[email protected]>
Co-authored-by: Michael Howitz <[email protected]>

Author: 3 people

docs/CHANGES.rst

@@ -4,6 +4,8 @@ Changes
 5.1a0 (unreleased)
 ------------------
 
+- Add support for (Python 3.8+) assignment expressions (i.e. the ``:=`` operator)
+
 - Fix ``compile_restricted_function`` with SyntaxErrors that have no text
   (`#181 <https://github.com/zopefoundation/RestrictedPython/issues/181>`_)
 

src/RestrictedPython/transformer.py

@@ -1473,3 +1473,22 @@ def visit_AsyncFor(self, node):
     def visit_AsyncWith(self, node):
         """Deny async functionality."""
         self.not_allowed(node)
+
+    # Assignment expressions (walrus operator ``:=``)
+    # New in 3.8
+    def visit_NamedExpr(self, node):
+        """Allow assignment expressions under some circumstances."""
+        # while the grammar requires ``node.target`` to be a ``Name``
+        # the abstract syntax is more permissive and allows an ``expr``.
+        # We support only a ``Name``.
+        # This is safe as the expression can only add/modify local
+        # variables. While this may hide global variables, an
+        # (implicitly performed) name check guarantees (as usual)
+        # that no essential global variable is hidden.
+        node = self.node_contents_visit(node)  # this checks ``node.target``
+        target = node.target
+        if not isinstance(target, ast.Name):
+            self.error(
+                node,
+                "Assignment expressions are only allowed for simple targets")
+        return node

tests/test_NamedExpr.py

@@ -0,0 +1,56 @@
+"""Assignment expression (``NamedExpr``) tests."""
+
+
+from ast import NodeTransformer
+from ast import parse
+from RestrictedPython import compile_restricted
+from RestrictedPython import safe_globals
+from RestrictedPython._compat import IS_PY38_OR_GREATER
+from unittest import skipUnless
+from unittest import TestCase
+
+
+@skipUnless(IS_PY38_OR_GREATER, "Feature available for Python 3.8+")
+class TestNamedExpr(TestCase):
+    def test_works(self):
+        code, gs = compile_str("if x:= x + 1: True\n")
+        gs["x"] = 0
+        exec(code, gs)
+        self.assertEqual(gs["x"], 1)
+
+    def test_no_private_target(self):
+        with self.assertRaises(SyntaxError):
+            compile_str("if _x_:= 1: True\n")
+
+    def test_simple_only(self):
+        # we test here that only a simple variable is allowed
+        # as assignemt expression target
+        # Currently (Python 3.8, 3.9), this is enforced by the
+        # Python concrete syntax; therefore, some (``ast``) trickery is
+        # necessary to produce a test for it.
+        class TransformNamedExprTarget(NodeTransformer):
+            def visit_NamedExpr(self, node):
+                # this is brutal but sufficient for the test
+                node.target = None
+                return node
+
+        mod = parse("if x:= x + 1: True\n")
+        mod = TransformNamedExprTarget().visit(mod)
+        with self.assertRaisesRegex(
+                SyntaxError,
+                "Assignment expressions are only allowed for simple target"):
+            code, gs = compile_str(mod)
+
+
+def compile_str(s, name="<unknown>"):
+    """code and globals for *s*.
+
+    *s* must be acceptable for ``compile_restricted`` (this is (especially) the
+    case for an ``str`` or ``ast.Module``).
+
+    *name* is a ``str`` used in error messages.
+    """
+    code = compile_restricted(s, name, 'exec')
+    gs = safe_globals.copy()
+    gs["__debug__"] = True  # assert active
+    return code, gs