From d7672bf7338640b47b28cd7d903fbab1a62d4f3d Mon Sep 17 00:00:00 2001
From: "Michael J. Sullivan" <sully@msully.net>
Date: Wed, 21 Oct 2020 16:22:56 -0700
Subject: [PATCH] [mypyc] Implement the walrus operator

My assessment that this would not be too hard was accurate.

Fixes mypyc/mypyc#765.
---
 mypyc/irbuild/expression.py       |  8 +++++
 mypyc/irbuild/visitor.py          |  5 ++--
 mypyc/test-data/run-python38.test | 50 +++++++++++++++++++++++++++++++
 mypyc/test/test_run.py            |  2 ++
 4 files changed, 62 insertions(+), 3 deletions(-)
 create mode 100644 mypyc/test-data/run-python38.test

diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py
index 2aec69fae34d..14c11e07090d 100644
--- 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
index cd6dec8890b3..67b8f04a7dc2 100644
--- 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
new file mode 100644
index 000000000000..beb553065f74
--- /dev/null
+++ 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
index 61d89caa08f7..82a288e0d293 100644
--- 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