Skip to content

Commit

Permalink
[mypyc] Implement the walrus operator (python#9624)
Browse files Browse the repository at this point in the history
My assessment that this would not be too hard was accurate.

Fixes mypyc/mypyc#765.
  • Loading branch information
msullivan authored Oct 22, 2020
1 parent ecec8bb commit 52b425b
Show file tree
Hide file tree
Showing 4 changed files with 62 additions and 3 deletions.
8 changes: 8 additions & 0 deletions mypyc/irbuild/expression.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
5 changes: 2 additions & 3 deletions mypyc/irbuild/visitor.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@
transform_dictionary_comprehension,
transform_slice_expr,
transform_generator_expr,
transform_assignment_expr,
)


Expand Down Expand Up @@ -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

Expand Down
50 changes: 50 additions & 0 deletions mypyc/test-data/run-python38.test
Original file line number Diff line number Diff line change
@@ -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([])) == []
2 changes: 2 additions & 0 deletions mypyc/test/test_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 52b425b

Please sign in to comment.