make basic things work

Author: JelleZijlstra

mypy/checkexpr.py

@@ -28,7 +28,7 @@
     OpExpr, UnaryExpr, IndexExpr, CastExpr, RevealExpr, TypeApplication, ListExpr,
     TupleExpr, DictExpr, LambdaExpr, SuperExpr, SliceExpr, Context, Expression,
     ListComprehension, GeneratorExpr, SetExpr, MypyFile, Decorator,
-    ConditionalExpr, ComparisonExpr, TempNode, SetComprehension,
+    ConditionalExpr, ComparisonExpr, TempNode, SetComprehension, AssignmentExpr,
     DictionaryComprehension, ComplexExpr, EllipsisExpr, StarExpr, AwaitExpr, YieldExpr,
     YieldFromExpr, TypedDictExpr, PromoteExpr, NewTypeExpr, NamedTupleExpr, TypeVarExpr,
     TypeAliasExpr, BackquoteExpr, EnumCallExpr, TypeAlias, SymbolNode, PlaceholderNode,
@@ -2421,6 +2421,11 @@ def check_list_multiply(self, e: OpExpr) -> Type:
         e.method_type = method_type
         return result
 
+    def visit_assignment_expr(self, e: AssignmentExpr) -> Type:
+        value = self.accept(e.value)
+        self.chk.check_assignment(e.target, e.value)
+        return value
+
     def visit_unary_expr(self, e: UnaryExpr) -> Type:
         """Type check an unary operation ('not', '-', '+' or '~')."""
         operand_type = self.accept(e.expr)

mypy/fastparse.py

@@ -22,7 +22,7 @@
     TupleExpr, GeneratorExpr, ListComprehension, ListExpr, ConditionalExpr,
     DictExpr, SetExpr, NameExpr, IntExpr, StrExpr, BytesExpr, UnicodeExpr,
     FloatExpr, CallExpr, SuperExpr, MemberExpr, IndexExpr, SliceExpr, OpExpr,
-    UnaryExpr, LambdaExpr, ComparisonExpr,
+    UnaryExpr, LambdaExpr, ComparisonExpr, AssignmentExpr,
     StarExpr, YieldFromExpr, NonlocalDecl, DictionaryComprehension,
     SetComprehension, ComplexExpr, EllipsisExpr, YieldExpr, Argument,
     AwaitExpr, TempNode, Expression, Statement,
@@ -718,10 +718,6 @@ def visit_AugAssign(self, n: ast3.AugAssign) -> OperatorAssignmentStmt:
                                    self.visit(n.value))
         return self.set_line(s, n)
 
-    def visit_NamedExpr(self, n: NamedExpr) -> None:
-        self.fail("assignment expressions are not yet supported", n.lineno, n.col_offset)
-        return None
-
     # For(expr target, expr iter, stmt* body, stmt* orelse, string? type_comment)
     def visit_For(self, n: ast3.For) -> ForStmt:
         if n.type_comment is not None:
@@ -883,6 +879,10 @@ def visit_Continue(self, n: ast3.Continue) -> ContinueStmt:
 
     # --- expr ---
 
+    def visit_NamedExpr(self, n: NamedExpr) -> AssignmentExpr:
+        s = AssignmentExpr(self.visit(n.target), self.visit(n.value))
+        return self.set_line(s, n)
+
     # BoolOp(boolop op, expr* values)
     def visit_BoolOp(self, n: ast3.BoolOp) -> OpExpr:
         # mypy translates (1 and 2 and 3) as (1 and (2 and 3))

mypy/literals.py

@@ -7,7 +7,7 @@
     ConditionalExpr, EllipsisExpr, YieldFromExpr, YieldExpr, RevealExpr, SuperExpr,
     TypeApplication, LambdaExpr, ListComprehension, SetComprehension, DictionaryComprehension,
     GeneratorExpr, BackquoteExpr, TypeVarExpr, TypeAliasExpr, NamedTupleExpr, EnumCallExpr,
-    TypedDictExpr, NewTypeExpr, PromoteExpr, AwaitExpr, TempNode,
+    TypedDictExpr, NewTypeExpr, PromoteExpr, AwaitExpr, TempNode, AssignmentExpr,
 )
 from mypy.visitor import ExpressionVisitor
 
@@ -159,6 +159,9 @@ def visit_index_expr(self, e: IndexExpr) -> Optional[Key]:
             return ('Index', literal_hash(e.base), literal_hash(e.index))
         return None
 
+    def visit_assignment_expr(self, e: AssignmentExpr) -> None:
+        return None
+
     def visit_call_expr(self, e: CallExpr) -> None:
         return None
 

mypy/newsemanal/semanal.py

@@ -73,7 +73,7 @@
     PlaceholderNode, COVARIANT, CONTRAVARIANT, INVARIANT,
     nongen_builtins, get_member_expr_fullname, REVEAL_TYPE,
     REVEAL_LOCALS, is_final_node, TypedDictExpr, type_aliases_target_versions,
-    EnumCallExpr
+    EnumCallExpr, AssignmentExpr,
 )
 from mypy.tvar_scope import TypeVarScope
 from mypy.typevars import fill_typevars
@@ -1832,6 +1832,10 @@ def visit_import_all(self, i: ImportAll) -> None:
     # Assignment
     #
 
+    def visit_assignment_expr(self, s: AssignmentExpr) -> None:
+        s.value.accept(self)
+        self.analyze_lvalue(s.target)
+
     def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         tag = self.track_incomplete_refs()
         s.rvalue.accept(self)

mypy/nodes.py

@@ -1581,6 +1581,17 @@ def accept(self, visitor: ExpressionVisitor[T]) -> T:
         return visitor.visit_unary_expr(self)
 
 
+class AssignmentExpr(Expression):
+    """Assignment expressions in Python 3.8+, like "a := 2"."""
+    def __init__(self, target: Expression, value: Expression) -> None:
+        super().__init__()
+        self.target = target
+        self.value = value
+
+    def accept(self, visitor: ExpressionVisitor[T]) -> T:
+        return visitor.visit_assignment_expr(self)
+
+
 # Map from binary operator id to related method name (in Python 3).
 op_methods = {
     '+': '__add__',

mypy/semanal.py

@@ -48,7 +48,7 @@
     GlobalDecl, SuperExpr, DictExpr, CallExpr, RefExpr, OpExpr, UnaryExpr,
     SliceExpr, CastExpr, RevealExpr, TypeApplication, Context, SymbolTable,
     SymbolTableNode, ListComprehension, GeneratorExpr,
-    LambdaExpr, MDEF, Decorator, SetExpr, TypeVarExpr,
+    LambdaExpr, MDEF, Decorator, SetExpr, TypeVarExpr, AssignmentExpr,
     StrExpr, BytesExpr, PrintStmt, ConditionalExpr, PromoteExpr,
     ComparisonExpr, StarExpr, ARG_POS, ARG_NAMED, type_aliases,
     YieldFromExpr, NamedTupleExpr, NonlocalDecl, SymbolNode,
@@ -1780,6 +1780,10 @@ def add_type_alias_deps(self, aliases_used: Iterable[str],
             target = self.scope.current_target()
         self.cur_mod_node.alias_deps[target].update(aliases_used)
 
+    def visit_assignment_expr(self, e: AssignmentExpr) -> None:
+        self.analyze_lvalue(e.target)
+        e.value.accept(self)
+
     def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         s.is_final_def = self.unwrap_final(s)
         self.analyze_lvalues(s)

mypy/semanal_pass1.py

@@ -154,6 +154,10 @@ def visit_block(self, b: Block) -> None:
             node.accept(self)
         self.sem.block_depth[-1] -= 1
 
+    def visit_assignment_expr(self, s: AssignmentStmt) -> None:
+        if self.sem.is_module_scope():
+            self.analyze_lvalue(s.target)
+
     def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         if self.sem.is_module_scope():
             for lval in s.lvalues:
@@ -333,6 +337,8 @@ def visit_decorator(self, d: Decorator) -> None:
 
     def visit_if_stmt(self, s: IfStmt) -> None:
         infer_reachability_of_if_statement(s, self.sem.options)
+        for expr in s.expr:
+            expr.accept(self)
         for node in s.body:
             node.accept(self)
         if s.else_body:

mypy/server/subexpr.py

@@ -7,6 +7,7 @@
     SliceExpr, CastExpr, RevealExpr, UnaryExpr, ListExpr, TupleExpr, DictExpr, SetExpr,
     IndexExpr, GeneratorExpr, ListComprehension, SetComprehension, DictionaryComprehension,
     ConditionalExpr, TypeApplication, LambdaExpr, StarExpr, BackquoteExpr, AwaitExpr,
+    AssignmentExpr,
 )
 from mypy.traverser import TraverserVisitor
 
@@ -102,6 +103,10 @@ def visit_reveal_expr(self, e: RevealExpr) -> None:
         self.add(e)
         super().visit_reveal_expr(e)
 
+    def visit_assignment_expr(self, e: AssignmentExpr) -> None:
+        self.add(e)
+        super().visit_assignment_expr(e)
+
     def visit_unary_expr(self, e: UnaryExpr) -> None:
         self.add(e)
         super().visit_unary_expr(e)

mypy/stats.py

@@ -15,7 +15,8 @@
 from mypy import nodes
 from mypy.nodes import (
     Expression, FuncDef, TypeApplication, AssignmentStmt, NameExpr, CallExpr, MypyFile,
-    MemberExpr, OpExpr, ComparisonExpr, IndexExpr, UnaryExpr, YieldFromExpr, RefExpr, ClassDef
+    MemberExpr, OpExpr, ComparisonExpr, IndexExpr, UnaryExpr, YieldFromExpr, RefExpr, ClassDef,
+    AssignmentExpr,
 )
 
 MYPY = False
@@ -167,6 +168,10 @@ def visit_index_expr(self, o: IndexExpr) -> None:
         self.process_node(o)
         super().visit_index_expr(o)
 
+    def visit_assignment_expr(self, o: AssignmentExpr) -> None:
+        self.process_node(o)
+        super().visit_assignment_expr(o)
+
     def visit_unary_expr(self, o: UnaryExpr) -> None:
         self.process_node(o)
         super().visit_unary_expr(o)

mypy/strconv.py

@@ -425,6 +425,9 @@ def visit_reveal_expr(self, o: 'mypy.nodes.RevealExpr') -> str:
             # REVEAL_LOCALS
             return self.dump([o.local_nodes], o)
 
+    def visit_assignment_expr(self, o: 'mypy.nodes.AssignmentExpr') -> str:
+        return self.dump([o.target, o.value], o)
+
     def visit_unary_expr(self, o: 'mypy.nodes.UnaryExpr') -> str:
         return self.dump([o.op, o.expr], o)
 

mypy/traverser.py

@@ -6,7 +6,7 @@
     ExpressionStmt, AssignmentStmt, OperatorAssignmentStmt, WhileStmt,
     ForStmt, ReturnStmt, AssertStmt, DelStmt, IfStmt, RaiseStmt,
     TryStmt, WithStmt, NameExpr, MemberExpr, OpExpr, SliceExpr, CastExpr, RevealExpr,
-    UnaryExpr, ListExpr, TupleExpr, DictExpr, SetExpr, IndexExpr,
+    UnaryExpr, ListExpr, TupleExpr, DictExpr, SetExpr, IndexExpr, AssignmentExpr,
     GeneratorExpr, ListComprehension, SetComprehension, DictionaryComprehension,
     ConditionalExpr, TypeApplication, ExecStmt, Import, ImportFrom,
     LambdaExpr, ComparisonExpr, OverloadedFuncDef, YieldFromExpr,
@@ -193,6 +193,10 @@ def visit_reveal_expr(self, o: RevealExpr) -> None:
             # RevealLocalsExpr doesn't have an inner expression
             pass
 
+    def visit_assignment_expr(self, o: AssignmentExpr) -> None:
+        o.target.accept(self)
+        o.value.accept(self)
+
     def visit_unary_expr(self, o: UnaryExpr) -> None:
         o.expr.accept(self)
 

mypy/treetransform.py

@@ -19,7 +19,7 @@
     ComparisonExpr, TempNode, StarExpr, Statement, Expression,
     YieldFromExpr, NamedTupleExpr, TypedDictExpr, NonlocalDecl, SetComprehension,
     DictionaryComprehension, ComplexExpr, TypeAliasExpr, EllipsisExpr,
-    YieldExpr, ExecStmt, Argument, BackquoteExpr, AwaitExpr,
+    YieldExpr, ExecStmt, Argument, BackquoteExpr, AwaitExpr, AssignmentExpr,
     OverloadPart, EnumCallExpr, REVEAL_TYPE
 )
 from mypy.types import Type, FunctionLike
@@ -407,6 +407,9 @@ def visit_super_expr(self, node: SuperExpr) -> SuperExpr:
         new.info = node.info
         return new
 
+    def visit_assignment_expr(self, node: AssignmentExpr) -> AssignmentExpr:
+        return AssignmentExpr(node.target, node.value)
+
     def visit_unary_expr(self, node: UnaryExpr) -> UnaryExpr:
         new = UnaryExpr(node.op, self.expr(node.expr))
         new.method_type = self.optional_type(node.method_type)

mypy/visitor.py

@@ -90,6 +90,10 @@ def visit_super_expr(self, o: 'mypy.nodes.SuperExpr') -> T:
     def visit_unary_expr(self, o: 'mypy.nodes.UnaryExpr') -> T:
         pass
 
+    @abstractmethod
+    def visit_assignment_expr(self, o: 'mypy.nodes.AssignmentExpr') -> T:
+        pass
+
     @abstractmethod
     def visit_list_expr(self, o: 'mypy.nodes.ListExpr') -> T:
         pass
@@ -473,6 +477,9 @@ def visit_reveal_expr(self, o: 'mypy.nodes.RevealExpr') -> T:
     def visit_super_expr(self, o: 'mypy.nodes.SuperExpr') -> T:
         pass
 
+    def visit_assignment_expr(self, o: 'mypy.nodes.AssignmentExpr') -> T:
+        pass
+
     def visit_unary_expr(self, o: 'mypy.nodes.UnaryExpr') -> T:
         pass
 

test-data/unit/check-38.test

@@ -106,3 +106,7 @@ def g(x: int): ...
     /
     0  # type: ignore
 )  # type: ignore  # E: unused 'type: ignore' comment
+
+[case testWalrus]
+if (a := 2):
+    reveal_type(a)  # E: Revealed type is 'builtins.int'