Support PEP 572

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,
@@ -720,10 +720,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:
@@ -885,6 +881,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
@@ -186,6 +186,8 @@ class NewSemanticAnalyzer(NodeVisitor[None],
     nonlocal_decls = None  # type: List[Set[str]]
     # Local names of function scopes; None for non-function scopes.
     locals = None  # type: List[Optional[SymbolTable]]
+    # Whether each scope is a comprehension scope.
+    is_comprehension_stack = None  # type: List[bool]
     # Nested block depths of scopes
     block_depth = None  # type: List[int]
     # TypeInfo of directly enclosing class (or None)
@@ -240,6 +242,7 @@ def __init__(self,
             errors: Report analysis errors using this instance
         """
         self.locals = [None]
+        self.is_comprehension_stack = [False]
         # Saved namespaces from previous iteration. Every top-level function/method body is
         # analyzed in several iterations until all names are resolved. We need to save
         # the local namespaces for the top level function and all nested functions between
@@ -510,9 +513,14 @@ def file_context(self,
     #
 
     def visit_func_def(self, defn: FuncDef) -> None:
+        # Visit default values because they may contain assignment expressions.
+        for arg in defn.arguments:
+            if arg.initializer:
+                arg.initializer.accept(self)
+
         defn.is_conditional = self.block_depth[-1] > 0
 
-        # Set full names even for those definitionss that aren't added
+        # Set full names even for those definitions that aren't added
         # to a symbol table. For example, for overload items.
         defn._fullname = self.qualified_name(defn.name())
 
@@ -1120,13 +1128,15 @@ def enter_class(self, info: TypeInfo) -> None:
         # Remember previous active class
         self.type_stack.append(self.type)
         self.locals.append(None)  # Add class scope
+        self.is_comprehension_stack.append(False)
         self.block_depth.append(-1)  # The class body increments this to 0
         self.type = info
 
     def leave_class(self) -> None:
         """ Restore analyzer state. """
         self.block_depth.pop()
         self.locals.pop()
+        self.is_comprehension_stack.pop()
         self.type = self.type_stack.pop()
 
     def analyze_class_decorator(self, defn: ClassDef, decorator: Expression) -> None:
@@ -1831,6 +1841,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, escape_comprehensions=True)
+
     def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         tag = self.track_incomplete_refs()
         s.rvalue.accept(self)
@@ -2396,16 +2410,22 @@ def check_and_set_up_type_alias(self, s: AssignmentStmt) -> bool:
 
     def analyze_lvalue(self, lval: Lvalue, nested: bool = False,
                        explicit_type: bool = False,
-                       is_final: bool = False) -> None:
+                       is_final: bool = False,
+                       escape_comprehensions: bool = False) -> None:
         """Analyze an lvalue or assignment target.
 
         Args:
             lval: The target lvalue
             nested: If true, the lvalue is within a tuple or list lvalue expression
             explicit_type: Assignment has type annotation
+            escape_comprehensions: If we are inside a comprehension, set the variable
+            in the enclosing scope instead. This implements
+            https://www.python.org/dev/peps/pep-0572/#scope-of-the-target
         """
+        if escape_comprehensions:
+            assert isinstance(lval, NameExpr), "assignment expression target must be NameExpr"
         if isinstance(lval, NameExpr):
-            self.analyze_name_lvalue(lval, explicit_type, is_final)
+            self.analyze_name_lvalue(lval, explicit_type, is_final, escape_comprehensions)
         elif isinstance(lval, MemberExpr):
             self.analyze_member_lvalue(lval, explicit_type, is_final)
             if explicit_type and not self.is_self_member_ref(lval):
@@ -2431,7 +2451,8 @@ def analyze_lvalue(self, lval: Lvalue, nested: bool = False,
     def analyze_name_lvalue(self,
                             lvalue: NameExpr,
                             explicit_type: bool,
-                            is_final: bool) -> None:
+                            is_final: bool,
+                            escape_comprehensions: bool) -> None:
         """Analyze an lvalue that targets a name expression.
 
         Arguments are similar to "analyze_lvalue".
@@ -2455,7 +2476,7 @@ def analyze_name_lvalue(self,
         if (not existing or isinstance(existing.node, PlaceholderNode)) and not outer:
             # Define new variable.
             var = self.make_name_lvalue_var(lvalue, kind, not explicit_type)
-            added = self.add_symbol(name, var, lvalue)
+            added = self.add_symbol(name, var, lvalue, escape_comprehensions=escape_comprehensions)
             # Only bind expression if we successfully added name to symbol table.
             if added:
                 lvalue.is_new_def = True
@@ -3948,7 +3969,8 @@ def add_symbol(self,
                    context: Context,
                    module_public: bool = True,
                    module_hidden: bool = False,
-                   can_defer: bool = True) -> bool:
+                   can_defer: bool = True,
+                   escape_comprehensions: bool = False) -> bool:
         """Add symbol to the currently active symbol table.
 
         Generally additions to symbol table should go through this method or
@@ -3970,7 +3992,7 @@ def add_symbol(self,
                                  node,
                                  module_public=module_public,
                                  module_hidden=module_hidden)
-        return self.add_symbol_table_node(name, symbol, context, can_defer)
+        return self.add_symbol_table_node(name, symbol, context, can_defer, escape_comprehensions)
 
     def add_symbol_skip_local(self, name: str, node: SymbolNode) -> None:
         """Same as above, but skipping the local namespace.
@@ -3998,7 +4020,8 @@ def add_symbol_table_node(self,
                               name: str,
                               symbol: SymbolTableNode,
                               context: Optional[Context] = None,
-                              can_defer: bool = True) -> bool:
+                              can_defer: bool = True,
+                              escape_comprehensions: bool = False) -> bool:
         """Add symbol table node to the currently active symbol table.
 
         Return True if we actually added the symbol, or False if we refused
@@ -4017,7 +4040,7 @@ def add_symbol_table_node(self,
             can_defer: if True, defer current target if adding a placeholder
             context: error context (see above about None value)
         """
-        names = self.current_symbol_table()
+        names = self.current_symbol_table(escape_comprehensions=escape_comprehensions)
         existing = names.get(name)
         if isinstance(symbol.node, PlaceholderNode) and can_defer:
             self.defer()
@@ -4250,13 +4273,16 @@ def enter(self, function: Union[FuncItem, GeneratorExpr, DictionaryComprehension
         """Enter a function, generator or comprehension scope."""
         names = self.saved_locals.setdefault(function, SymbolTable())
         self.locals.append(names)
+        is_comprehension = isinstance(function, (GeneratorExpr, DictionaryComprehension))
+        self.is_comprehension_stack.append(is_comprehension)
         self.global_decls.append(set())
         self.nonlocal_decls.append(set())
         # -1 since entering block will increment this to 0.
         self.block_depth.append(-1)
 
     def leave(self) -> None:
         self.locals.pop()
+        self.is_comprehension_stack.pop()
         self.global_decls.pop()
         self.nonlocal_decls.pop()
         self.block_depth.pop()
@@ -4283,10 +4309,19 @@ def current_symbol_kind(self) -> int:
             kind = GDEF
         return kind
 
-    def current_symbol_table(self) -> SymbolTable:
+    def current_symbol_table(self, escape_comprehensions: bool = False) -> SymbolTable:
         if self.is_func_scope():
             assert self.locals[-1] is not None
-            names = self.locals[-1]
+            if escape_comprehensions:
+                for i, is_comprehension in enumerate(reversed(self.is_comprehension_stack)):
+                    if not is_comprehension:
+                        names = self.locals[-1 - i]
+                        break
+                else:
+                    assert False, "Should have at least one non-comprehension scope"
+            else:
+                names = self.locals[-1]
+            assert names is not None
         elif self.type is not None:
             names = self.type.names
         else:

mypy/newsemanal/semanal_pass1.py

@@ -90,6 +90,8 @@ def visit_import(self, node: Import) -> None:
 
     def visit_if_stmt(self, s: IfStmt) -> None:
         infer_reachability_of_if_statement(s, self.options)
+        for expr in s.expr:
+            expr.accept(self)
         for node in s.body:
             node.accept(self)
         if s.else_body:

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__',