Support PEP 572 (#6899)

Author: JelleZijlstra

mypy/checker.py

@@ -24,7 +24,7 @@
     ComparisonExpr, StarExpr, EllipsisExpr, RefExpr, PromoteExpr,
     Import, ImportFrom, ImportAll, ImportBase, TypeAlias,
     ARG_POS, ARG_STAR, LITERAL_TYPE, MDEF, GDEF,
-    CONTRAVARIANT, COVARIANT, INVARIANT, TypeVarExpr,
+    CONTRAVARIANT, COVARIANT, INVARIANT, TypeVarExpr, AssignmentExpr,
     is_final_node,
 )
 from mypy import nodes
@@ -2225,7 +2225,8 @@ def enter_final_context(self, is_final_def: bool) -> Iterator[None]:
         finally:
             self._is_final_def = old_ctx
 
-    def check_final(self, s: Union[AssignmentStmt, OperatorAssignmentStmt]) -> None:
+    def check_final(self,
+                    s: Union[AssignmentStmt, OperatorAssignmentStmt, AssignmentExpr]) -> None:
         """Check if this assignment does not assign to a final attribute.
 
         This function performs the check only for name assignments at module
@@ -2234,6 +2235,8 @@ def check_final(self, s: Union[AssignmentStmt, OperatorAssignmentStmt]) -> None:
         """
         if isinstance(s, AssignmentStmt):
             lvs = self.flatten_lvalues(s.lvalues)
+        elif isinstance(s, AssignmentExpr):
+            lvs = [s.target]
         else:
             lvs = [s.lvalue]
         is_final_decl = s.is_final_def if isinstance(s, AssignmentStmt) else False

mypy/checkexpr.py

@@ -26,7 +26,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,
@@ -2544,6 +2544,12 @@ 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)
+        self.chk.check_final(e)
+        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

@@ -20,7 +20,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,
@@ -757,10 +757,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:
         target_type = self.translate_type_comment(n, n.type_comment)
@@ -902,6 +898,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

@@ -8,7 +8,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
 
@@ -156,6 +156,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/nodes.py

@@ -1593,6 +1593,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

@@ -74,7 +74,7 @@
     PlaceholderNode, COVARIANT, CONTRAVARIANT, INVARIANT,
     nongen_builtins, get_member_expr_fullname, REVEAL_TYPE,
     REVEAL_LOCALS, is_final_node, TypedDictExpr, type_aliases_target_versions,
-    EnumCallExpr, RUNTIME_PROTOCOL_DECOS, FakeExpression, Statement
+    EnumCallExpr, RUNTIME_PROTOCOL_DECOS, FakeExpression, Statement, AssignmentExpr,
 )
 from mypy.tvar_scope import TypeVarScope
 from mypy.typevars import fill_typevars
@@ -177,6 +177,8 @@ class SemanticAnalyzer(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)
@@ -242,6 +244,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
@@ -519,6 +522,12 @@ def file_context(self,
 
     def visit_func_def(self, defn: FuncDef) -> None:
         self.statement = defn
+
+        # 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 definitions that aren't added
@@ -1148,13 +1157,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:
@@ -1858,6 +1869,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:
         self.statement = s
 
@@ -2493,16 +2508,22 @@ 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):
@@ -2528,7 +2549,8 @@ def analyze_lvalue(self,
     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".
@@ -2552,7 +2574,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
@@ -4082,7 +4104,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
@@ -4104,7 +4127,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.
@@ -4132,7 +4155,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
@@ -4151,7 +4175,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(context)
@@ -4379,13 +4403,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()
@@ -4412,10 +4439,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/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/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

@@ -18,8 +18,8 @@
 from mypy.nodes import (
     Expression, FuncDef, TypeApplication, AssignmentStmt, NameExpr, CallExpr, MypyFile,
     MemberExpr, OpExpr, ComparisonExpr, IndexExpr, UnaryExpr, YieldFromExpr, RefExpr, ClassDef,
-    ImportFrom, Import, ImportAll, PassStmt, BreakStmt, ContinueStmt, StrExpr, BytesExpr,
-    UnicodeExpr, IntExpr, FloatExpr, ComplexExpr, EllipsisExpr, ExpressionStmt, Node
+    AssignmentExpr, ImportFrom, Import, ImportAll, PassStmt, BreakStmt, ContinueStmt, StrExpr,
+    BytesExpr, UnicodeExpr, IntExpr, FloatExpr, ComplexExpr, EllipsisExpr, ExpressionStmt, Node
 )
 from mypy.util import correct_relative_import
 from mypy.argmap import map_formals_to_actuals
@@ -275,6 +275,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,
@@ -192,6 +192,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)