Add support for PEP 572 - Assignment expressions (#134)

Closes #126

astor/code_gen.py

@@ -741,6 +741,22 @@ class SourceGenerator(ExplicitNodeVisitor):
             for op, right in zip(node.ops, node.comparators):
                 self.write(get_op_symbol(op, ' %s '), right)
 
+    # assignment expressions; new for Python 3.8
+    def visit_NamedExpr(self, node):
+        with self.delimit(node) as delimiters:
+            p = delimiters.p
+            set_precedence(p, node.target)
+            set_precedence(p + 1, node.value)
+            # Python is picky about delimiters for assignment
+            # expressions: it requires at least one pair in any
+            # statement that uses an assignment expression, even
+            # when not necessary according to the precedence
+            # rules. We address this with the kludge of forcing a
+            # pair of parentheses around every assignment
+            # expression.
+            delimiters.discard = False
+            self.write(node.target, ' := ', node.value)
+
     def visit_UnaryOp(self, node):
         with self.delimit(node, node.op) as delimiters:
             set_precedence(delimiters.p, node.operand)

astor/op_util.py

@@ -36,6 +36,7 @@ op_data = """
            Tuple                0
 
            Comma                1
+       NamedExpr                1
           Assert                0
            Raise                0
     call_one_arg                1

docs/changelog.rst

@@ -15,6 +15,12 @@ New features
 .. _`Issue 120`: https://github.com/berkerpeksag/astor/issues/120
 .. _`PR 121`: https://github.com/berkerpeksag/astor/pull/121
 
+* Support Python 3.8's assignment expressions.
+  (Reported and fixed by Kodi Arfer in `Issue 126`_ and `PR 134`_.)
+
+.. _`Issue 126`: https://github.com/berkerpeksag/astor/issues/126
+.. _`PR 134`: https://github.com/berkerpeksag/astor/pull/134
+
 Bug fixes
 ~~~~~~~~~
 

tests/test_code_gen.py

@@ -454,6 +454,27 @@ class CodegenTestCase(unittest.TestCase, Comparisons):
         """
         self.assertSrcRoundtripsGtVer(source, (3, 6))
 
+    def test_assignment_expr(self):
+        cases = (
+            "(x := 3)",
+            "1 + (x := y)",
+            "x = (y := 0)",
+            "1 + (p := 1 if 2 else 3)",
+            "[y := f(x), y**2, y**3]",
+            "(2 ** 3 * 4 + 5 and 6, x := 2 ** 3 * 4 + 5 and 6)",
+            "foo(x := 3, cat='vector')",
+            "foo(x=(y := f(x)))",
+            "any(len(longline := line) >= 100 for line in lines)",
+            "[[y := f(x), x/y] for x in range(5)]",
+            "lambda: (x := 1)",
+            "def foo(answer=(p := 42)): pass",
+            "def foo(answer: (p := 42) = 5): pass",
+            "if reductor := dispatch_table.get(cls): pass",
+            "while line := fp.readline(): pass",
+            "while (command := input('> ')) != 'quit': pass")
+        for case in cases:
+            self.assertAstRoundtripsGtVer(case, (3, 8))
+
     @unittest.skipUnless(sys.version_info <= (3, 3),
                          "ast.Name used for True, False, None until Python 3.4")
     def test_deprecated_constants_as_name(self):