Implement new NamedTuple syntax

gvanrossum · Oct 25, 2016 · Oct 13, 2016 · Oct 13, 2016 · Oct 13, 2016 · Oct 16, 2016
commit 409abb668c5a2966a15790bfa7867b725e2a895f
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -53,7 +53,7 @@
     ImportFrom, ImportAll, Block, LDEF, NameExpr, MemberExpr,
     IndexExpr, TupleExpr, ListExpr, ExpressionStmt, ReturnStmt,
     RaiseStmt, AssertStmt, OperatorAssignmentStmt, WhileStmt,
-    ForStmt, BreakStmt, ContinueStmt, IfStmt, TryStmt, WithStmt, DelStmt,
+    ForStmt, BreakStmt, ContinueStmt, IfStmt, TryStmt, WithStmt, DelStmt, PassStmt,
     GlobalDecl, SuperExpr, DictExpr, CallExpr, RefExpr, OpExpr, UnaryExpr,
     SliceExpr, CastExpr, RevealTypeExpr, TypeApplication, Context, SymbolTable,
     SymbolTableNode, BOUND_TVAR, UNBOUND_TVAR, ListComprehension, GeneratorExpr,
@@ -63,7 +63,7 @@
     YieldFromExpr, NamedTupleExpr, NonlocalDecl, SymbolNode,
     SetComprehension, DictionaryComprehension, TYPE_ALIAS, TypeAliasExpr,
     YieldExpr, ExecStmt, Argument, BackquoteExpr, ImportBase, AwaitExpr,
-    IntExpr, FloatExpr, UnicodeExpr, EllipsisExpr,
+    IntExpr, FloatExpr, UnicodeExpr, EllipsisExpr, TempNode,
     COVARIANT, CONTRAVARIANT, INVARIANT, UNBOUND_IMPORTED, LITERAL_YES,
 )
 from mypy.visitor import NodeVisitor
@@ -545,7 +545,54 @@ def check_function_signature(self, fdef: FuncItem) -> None:
         elif len(sig.arg_types) > len(fdef.arguments):
             self.fail('Type signature has too many arguments', fdef, blocker=True)
 
+    def check_namedtuple_classdef(self, defn: ClassDef) -> Tuple[List[str], List[Type]]:
+        if self.options.python_version < (3, 6):
+            self.fail('NamedTuple class syntax is only supported in Python 3.6', defn)
+            return [], []
+        if len(defn.base_type_exprs) > 1:
+            self.fail('NamedTuple should be a single base', defn)
+            return [], []
+        items = []  # type: List[str]
+        types = []  # type: List[Type]
+        for stmt in defn.defs.body:
+            if isinstance(stmt, AssignmentStmt):
+                if len(stmt.lvalues) > 1 or not isinstance(stmt.lvalues[0], NameExpr):
+                    self.fail('NamedTuple supports only simple fields', stmt)
+                    return [], []
+                name = stmt.lvalues[0].name
+                if not isinstance(stmt.rvalue, TempNode):
+                    # x: int assigns rvalue to TempNode(AnyType())
+                    self.fail('NamedTuple does not support initial values', stmt)
+                if stmt.type is None:
+                    self.fail('NamedTuple field type must be specified', stmt)
+                    types.append(AnyType())
+                else:
+                    types.append(self.anal_type(stmt.type))
+                if name.startswith('_'):
+                    self.fail('NamedTuple field name cannot start with an underscore: {}'
+                              .format(name), stmt)
+                items.append(name)
+            else:
+                if not isinstance(stmt, PassStmt):
+                    self.fail('Invalid definition in NamedTuple', stmt)
+        return items, types
+
+    def analyze_namedtuple_classdef(self, defn: ClassDef) -> None:
+        node = self.lookup(defn.name, defn)
+        node.kind = GDEF  # TODO in process_namedtuple_definition also applies here
+        items, types = self.check_namedtuple_classdef(defn)
+        node.node = self.build_namedtuple_typeinfo(defn.name, items, types)
+
     def visit_class_def(self, defn: ClassDef) -> None:
+        # special case for NamedTuple
+        for base_expr in defn.base_type_exprs:
+            if isinstance(base_expr, NameExpr):
+                sym = self.lookup_qualified(base_expr.name, base_expr)
+                if sym is None or sym.node is None:
+                    continue
+                if sym.node.fullname() == 'typing.NamedTuple':
+                    self.analyze_namedtuple_classdef(defn)
+                    return
         self.clean_up_bases_and_infer_type_variables(defn)
         self.setup_class_def_analysis(defn)
 

diff --git a/test-data/unit/check-class-namedtuple.test b/test-data/unit/check-class-namedtuple.test
@@ -2,18 +2,17 @@
 # flags: --fast-parser --python-version 3.5
 from typing import NamedTuple
 
-class E(NamedTuple):  # E: NamedTuple class syntax is only supported in Python 3.6+
-    x: int
-    y: str
+class E(NamedTuple):  # E: NamedTuple class syntax is only supported in Python 3.6
+    pass
 
 [case testNewNamedTupleNoUnderscoreFields]
 # flags: --fast-parser --python-version 3.6
 from typing import NamedTuple
 
-class X(NamedTuple):  # E: NamedTuple field names cannot start with an underscore: _y, _z
+class X(NamedTuple):
     x: int
-    _y: int
-    _z: int
+    _y: int  # E: NamedTuple field name cannot start with an underscore: _y
+    _z: int  # E: NamedTuple field name cannot start with an underscore: _z
 
 [case testNewNamedTupleAccessingAttributes]
 # flags: --fast-parser --python-version 3.6
@@ -34,18 +33,14 @@ from typing import NamedTuple
 
 class X(NamedTuple):
     x: int
-    y: str
 
 x: X
 x.x = 5 # E: Property "x" defined in "X" is read-only
-x.y = 5 # E: Property "y" defined in "X" is read-only
-x.z = 5 # E: "X" has no attribute "z"
+x.y = 5 # E: "X" has no attribute "y"
 
 class A(X): pass
 a: A
 a.x = 5 # E: Property "x" defined in "A" is read-only
-a.y = 5 # E: Property "y" defined in "A" is read-only
--- a.z = 5 # not supported yet
 
 [case testNewNamedTupleCreateWithPositionalArguments]
 # flags: --fast-parser --python-version 3.6
@@ -55,11 +50,11 @@ class X(NamedTuple):
     x: int
     y: str
 
-x = X(1, 'x')
+x = X(1, '2')
 x.x
 x.z      # E: "X" has no attribute "z"
 x = X(1) # E: Too few arguments for "X"
-x = X(1, 2, 3)  # E: Too many arguments for "X"
+x = X(1, '2', 3)  # E: Too many arguments for "X"
 
 [case testCreateNewNamedTupleWithKeywordArguments]
 # flags: --fast-parser --python-version 3.6
@@ -72,7 +67,7 @@ class X(NamedTuple):
 x = X(x=1, y='x')
 x = X(1, y='x')
 x = X(x=1, z=1) # E: Unexpected keyword argument "z" for "X"
-x = X(y=1) # E: Missing positional argument "x" in call to "X"
+x = X(y='x') # E: Missing positional argument "x" in call to "X"
 
 [case testNewNamedTupleCreateAndUseAsTuple]
 # flags: --fast-parser --python-version 3.6
@@ -227,15 +222,14 @@ MyNamedTuple.x # E: "MyNamedTuple" has no attribute "x"
 from typing import NamedTuple
 
 class A(NamedTuple):
-    ...
+    pass
 
 [case testNewNamedTupleProperty]
 # flags: --fast-parser --python-version 3.6
 from typing import NamedTuple
 
 class A(NamedTuple):
     a: int
-    b: str
 
 class B(A):
     @property
@@ -288,7 +282,7 @@ reveal_type(X._fields)  # E: Revealed type is 'Tuple[builtins.str, builtins.str]
 from typing import NamedTuple
 
 class X(NamedTuple):
-    ...
+    pass
 
 x: X = X()
 x._replace()
@@ -329,7 +323,7 @@ from typing import NamedTuple
 class X(NamedTuple):
     x: int
     y: str
-    def f(self): pass  # E: NamedTuple only allows names as fields
+    def f(self): pass  # E: Invalid definition in NamedTuple
 
 [case testNewNamedTupleWithInvalidItems2]
 # flags: --fast-parser --python-version 3.6
@@ -338,13 +332,13 @@ from typing import NamedTuple
 class X(NamedTuple):
     x: int
     y: str
-    [0][0]: int  # E: NamedTuple only allows names as fields
+    x[0]: int  # E: NamedTuple supports only simple fields
 
 [builtins fixtures/list.pyi]
 
 [case testTypeUsingTypeCNamedTuple]
 # flags: --fast-parser --python-version 3.6
-from typing import NamedTuple
+from typing import NamedTuple, Type
 
 class N(NamedTuple):
     x: int
@@ -355,4 +349,4 @@ def f(a: Type[N]):
 [builtins fixtures/list.pyi]
 [out]
 main: note: In function "f":
-main:3: error: Unsupported type Type["N"]
+main:8: error: Unsupported type Type["N"]