Implement PEP 526 Variable Annotations Syntax (#2131)

Depends on python/typed_ast#16 for the new syntax, but is backwards compatible otherwise.

Author: ilevkivskyi

mypy/checker.py

@@ -1030,7 +1030,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> Type:
 
         Handle all kinds of assignment statements (simple, indexed, multiple).
         """
-        self.check_assignment(s.lvalues[-1], s.rvalue, s.type is None)
+        self.check_assignment(s.lvalues[-1], s.rvalue, s.type is None, s.new_syntax)
 
         if len(s.lvalues) > 1:
             # Chained assignment (e.g. x = y = ...).
@@ -1041,7 +1041,8 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> Type:
             for lv in s.lvalues[:-1]:
                 self.check_assignment(lv, rvalue, s.type is None)
 
-    def check_assignment(self, lvalue: Node, rvalue: Node, infer_lvalue_type: bool = True) -> None:
+    def check_assignment(self, lvalue: Node, rvalue: Node, infer_lvalue_type: bool = True,
+                         new_syntax: bool = False) -> None:
         """Type check a single assignment: lvalue = rvalue."""
         if isinstance(lvalue, TupleExpr) or isinstance(lvalue, ListExpr):
             self.check_assignment_to_multiple_lvalues(lvalue.items, rvalue, lvalue,
@@ -1078,7 +1079,8 @@ def check_assignment(self, lvalue: Node, rvalue: Node, infer_lvalue_type: bool =
                 elif (is_literal_none(rvalue) and
                         isinstance(lvalue, NameExpr) and
                         isinstance(lvalue.node, Var) and
-                        lvalue.node.is_initialized_in_class):
+                        lvalue.node.is_initialized_in_class and
+                        not new_syntax):
                     # Allow None's to be assigned to class variables with non-Optional types.
                     rvalue_type = lvalue_type
                 else:

mypy/fastparse.py

@@ -14,7 +14,7 @@
     UnaryExpr, FuncExpr, ComparisonExpr,
     StarExpr, YieldFromExpr, NonlocalDecl, DictionaryComprehension,
     SetComprehension, ComplexExpr, EllipsisExpr, YieldExpr, Argument,
-    AwaitExpr,
+    AwaitExpr, TempNode,
     ARG_POS, ARG_OPT, ARG_STAR, ARG_NAMED, ARG_STAR2
 )
 from mypy.types import (
@@ -403,16 +403,31 @@ def visit_Delete(self, n: ast35.Delete) -> Node:
         else:
             return DelStmt(self.visit(n.targets[0]))
 
-    # Assign(expr* targets, expr value, string? type_comment)
+    # Assign(expr* targets, expr? value, string? type_comment, expr? annotation)
     @with_line
     def visit_Assign(self, n: ast35.Assign) -> Node:
         typ = None
-        if n.type_comment:
+        if hasattr(n, 'annotation') and n.annotation is not None:  # type: ignore
+            new_syntax = True
+        else:
+            new_syntax = False
+        if new_syntax and self.pyversion < (3, 6):
+            raise TypeCommentParseError('Variable annotation syntax is only '
+                                        'suppoted in Python 3.6, use type '
+                                        'comment instead', n.lineno, n.col_offset)
+        # typed_ast prevents having both type_comment and annotation.
+        if n.type_comment is not None:
             typ = parse_type_comment(n.type_comment, n.lineno)
-
-        return AssignmentStmt(self.visit_list(n.targets),
-                              self.visit(n.value),
-                              type=typ)
+        elif new_syntax:
+            typ = TypeConverter(line=n.lineno).visit(n.annotation)  # type: ignore
+        if n.value is None:  # always allow 'x: int'
+            rvalue = TempNode(AnyType())  # type: Node
+        else:
+            rvalue = self.visit(n.value)
+        lvalues = self.visit_list(n.targets)
+        return AssignmentStmt(lvalues,
+                              rvalue,
+                              type=typ, new_syntax=new_syntax)
 
     # AugAssign(expr target, operator op, expr value)
     @with_line

mypy/nodes.py

@@ -779,12 +779,15 @@ class AssignmentStmt(Statement):
     rvalue = None  # type: Expression
     # Declared type in a comment, may be None.
     type = None  # type: mypy.types.Type
+    # This indicates usage of PEP 526 type annotation syntax in assignment.
+    new_syntax = False  # type: bool
 
     def __init__(self, lvalues: List[Expression], rvalue: Expression,
-                 type: 'mypy.types.Type' = None) -> None:
+                 type: 'mypy.types.Type' = None, new_syntax: bool = False) -> None:
         self.lvalues = lvalues
         self.rvalue = rvalue
         self.type = type
+        self.new_syntax = new_syntax
 
     def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit_assignment_stmt(self)
@@ -1786,6 +1789,9 @@ class TempNode(Expression):
     def __init__(self, typ: 'mypy.types.Type') -> None:
         self.type = typ
 
+    def __repr__(self):
+        return 'TempNode(%s)' % str(self.type)
+
     def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit_temp_node(self)
 

mypy/test/testcheck.py

@@ -5,6 +5,8 @@
 import shutil
 import sys
 import time
+import typed_ast
+import typed_ast.ast35
 
 from typing import Tuple, List, Dict, Set
 
@@ -68,6 +70,9 @@
     'check-columns.test',
 ]
 
+if 'annotation' in typed_ast.ast35.Assign._fields:
+    files.append('check-newsyntax.test')
+
 
 class TypeCheckSuite(DataSuite):
     def __init__(self, *, update_data=False):

test-data/unit/check-newsyntax.test

@@ -0,0 +1,90 @@
+[case testNewSyntaxRequire36]
+# flags: --fast-parser --python-version 3.5
+x: int = 5  # E: Variable annotation syntax is only suppoted in Python 3.6, use type comment instead
+[out]
+
+[case testNewSyntaxSyntaxError]
+# flags: --fast-parser --python-version 3.6
+x: int: int  # E: invalid syntax
+[out]
+
+[case testNewSyntaxBasics]
+# flags: --fast-parser --python-version 3.6
+x: int
+x = 5
+y: int = 5
+
+a: str
+a = 5  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
+b: str = 5  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
+
+zzz: int
+zzz: str  # E: Name 'zzz' already defined
+[out]
+
+[case testNewSyntaxWithDict]
+# flags: --fast-parser --python-version 3.6
+from typing import Dict, Any
+
+d: Dict[int, str] = {}
+d[42] = 'ab'
+d[42] = 42  # E: Incompatible types in assignment (expression has type "int", target has type "str")
+d['ab'] = 'ab'  # E: Invalid index type "str" for "dict"
+[builtins fixtures/dict.pyi]
+[out]
+
+[case testNewSyntaxWithRevealType]
+# flags: --fast-parser --python-version 3.6
+from typing import Dict
+
+def tst_local(dct: Dict[int, T]) -> Dict[T, int]:
+    ret: Dict[T, int] = {}
+    return ret
+
+reveal_type(tst_local({1: 'a'}))  # E: Revealed type is 'builtins.dict[builtins.str*, builtins.int]'
+[builtins fixtures/dict.pyi]
+[out]
+
+[case testNewSyntaxWithInstanceVars]
+# flags: --fast-parser --python-version 3.6
+class TstInstance:
+    a: str
+    def __init__(self) -> None:
+        self.x: int
+
+TstInstance().x = 5
+TstInstance().x = 'ab'  # E: Incompatible types in assignment (expression has type "str", variable has type "int")
+TstInstance().a = 5  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
+TstInstance().a = 'ab'
+[out]
+
+[case testNewSyntaxWithClassVars]
+# flags: --fast-parser --strict-optional --python-version 3.6
+class CCC:
+    a: str = None  # E: Incompatible types in assignment (expression has type None, variable has type "str")
+[out]
+main: note: In class "CCC":
+
+[case testNewSyntaxWithStrictOptional]
+# flags: --fast-parser --strict-optional --python-version 3.6
+strict: int
+strict = None  # E: Incompatible types in assignment (expression has type None, variable has type "int")
+strict2: int = None  # E: Incompatible types in assignment (expression has type None, variable has type "int")
+[out]
+
+[case testNewSyntaxWithStrictOptionalFunctions]
+# flags: --fast-parser --strict-optional --python-version 3.6
+def f() -> None:
+    x: int
+    x = None  # E: Incompatible types in assignment (expression has type None, variable has type "int")
+[out]
+main: note: In function "f":
+
+[case testNewSyntaxWithStrictOptionalClasses]
+# flags: --fast-parser --strict-optional --python-version 3.6
+class C:
+    def meth(self) -> None:
+        x: int = None  # E: Incompatible types in assignment (expression has type None, variable has type "int")
+        self.x: int = None  # E: Incompatible types in assignment (expression has type None, variable has type "int")
+[out]
+main: note: In member "meth" of class "C":