[mypyc] Implement the walrus operator (python#9624)

msullivan · web-flow · commit 52b425b00721 · 2020-10-22T11:12:36.000+08:00
My assessment that this would not be too hard was accurate. Fixes mypyc/mypyc#765.
diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py
@@ -11,6 +11,7 @@
     ConditionalExpr, ComparisonExpr, IntExpr, FloatExpr, ComplexExpr, StrExpr,
     BytesExpr, EllipsisExpr, ListExpr, TupleExpr, DictExpr, SetExpr, ListComprehension,
     SetComprehension, DictionaryComprehension, SliceExpr, GeneratorExpr, CastExpr, StarExpr,
+    AssignmentExpr,
     Var, RefExpr, MypyFile, TypeInfo, TypeApplication, LDEF, ARG_POS
 )
 from mypy.types import TupleType, get_proper_type, Instance
@@ -682,3 +683,10 @@ def transform_generator_expr(builder: IRBuilder, o: GeneratorExpr) -> Value:
     return builder.call_c(
         iter_op, [translate_list_comprehension(builder, o)], o.line
     )
+
+
+def transform_assignment_expr(builder: IRBuilder, o: AssignmentExpr) -> Value:
+    value = builder.accept(o.value)
+    target = builder.get_assignment_target(o.target)
+    builder.assign(target, value, o.line)
+    return value
diff --git a/mypyc/irbuild/visitor.py b/mypyc/irbuild/visitor.py
@@ -76,6 +76,7 @@
     transform_dictionary_comprehension,
     transform_slice_expr,
     transform_generator_expr,
+    transform_assignment_expr,
 )
 
 
@@ -264,10 +265,8 @@ def visit_yield_from_expr(self, o: YieldFromExpr) -> Value:
     def visit_await_expr(self, o: AwaitExpr) -> Value:
         return transform_await_expr(self.builder, o)
 
-    # Unimplemented constructs
-
     def visit_assignment_expr(self, o: AssignmentExpr) -> Value:
-        self.bail("I Am The Walrus (unimplemented)", o.line)
+        return transform_assignment_expr(self.builder, o)
 
     # Unimplemented constructs that shouldn't come up because they are py2 only
 
diff --git a/mypyc/test-data/run-python38.test b/mypyc/test-data/run-python38.test
@@ -0,0 +1,50 @@
+[case testWalrus1]
+from typing import Optional
+
+def foo(x: int) -> Optional[int]:
+    if x < 0:
+        return None
+    return x
+
+def test(x: int) -> str:
+    if (n := foo(x)) is not None:
+        return str(x)
+    else:
+        return "<fail>"
+
+[file driver.py]
+from native import test
+
+assert test(10) == "10"
+assert test(-1) == "<fail>"
+
+
+[case testWalrus2]
+from typing import Optional, Tuple, List
+
+class Node:
+    def __init__(self, val: int, next: Optional['Node']) -> None:
+        self.val = val
+        self.next = next
+
+def pairs(nobe: Optional[Node]) -> List[Tuple[int, int]]:
+    if nobe is None:
+        return []
+    l = []
+    while next := nobe.next:
+        l.append((nobe.val, next.val))
+        nobe = next
+    return l
+
+def make(l: List[int]) -> Optional[Node]:
+    cur: Optional[Node] = None
+    for x in reversed(l):
+        cur = Node(x, cur)
+    return cur
+
+[file driver.py]
+from native import Node, make, pairs
+
+assert pairs(make([1,2,3])) == [(1,2), (2,3)]
+assert pairs(make([1])) == []
+assert pairs(make([])) == []
diff --git a/mypyc/test/test_run.py b/mypyc/test/test_run.py
@@ -50,6 +50,8 @@
     'run-bench.test',
     'run-mypy-sim.test',
 ]
+if sys.version_info >= (3, 8):
+    files.append('run-python38.test')
 
 setup_format = """\
 from setuptools import setup

Original file line number	Diff line number	Diff line change
`@@ -50,6 +50,8 @@`
`50`	`50`	`'run-bench.test',`
`51`	`51`	`'run-mypy-sim.test',`
`52`	`52`	`]`
	`53`	`+if sys.version_info >= (3, 8):`
	`54`	`+ files.append('run-python38.test')`
`53`	`55`
`54`	`56`	`setup_format = """\`
`55`	`57`	`from setuptools import setup`