Implement default values on NamedTuple (#2719)

This implements python/typing#338 in mypy.

Author: JelleZijlstra

mypy/nodes.py

@@ -427,7 +427,7 @@ class Argument(Node):
     variable = None  # type: Var
     type_annotation = None  # type: Optional[mypy.types.Type]
     initializer = None  # type: Optional[Expression]
-    kind = None  # type: int
+    kind = None  # type: int  # must be an ARG_* constant
     initialization_statement = None  # type: Optional[AssignmentStmt]
 
     def __init__(self, variable: 'Var', type_annotation: 'Optional[mypy.types.Type]',

mypy/semanal.py

@@ -65,7 +65,7 @@
     SetComprehension, DictionaryComprehension, TYPE_ALIAS, TypeAliasExpr,
     YieldExpr, ExecStmt, Argument, BackquoteExpr, ImportBase, AwaitExpr,
     IntExpr, FloatExpr, UnicodeExpr, EllipsisExpr, TempNode,
-    COVARIANT, CONTRAVARIANT, INVARIANT, UNBOUND_IMPORTED, LITERAL_YES, nongen_builtins,
+    COVARIANT, CONTRAVARIANT, INVARIANT, UNBOUND_IMPORTED, LITERAL_YES, ARG_OPT, nongen_builtins,
     collections_type_aliases, get_member_expr_fullname,
 )
 from mypy.typevars import has_no_typevars, fill_typevars
@@ -588,27 +588,32 @@ def visit_class_def(self, defn: ClassDef) -> None:
         if self.analyze_typeddict_classdef(defn):
             return
         if self.analyze_namedtuple_classdef(defn):
-            return
-        self.setup_class_def_analysis(defn)
+            # just analyze the class body so we catch type errors in default values
+            self.enter_class(defn)
+            defn.defs.accept(self)
+            self.leave_class()
+        else:
+            self.setup_class_def_analysis(defn)
 
-        self.bind_class_type_vars(defn)
+            self.bind_class_type_vars(defn)
 
-        self.analyze_base_classes(defn)
-        self.analyze_metaclass(defn)
+            self.analyze_base_classes(defn)
+            self.analyze_metaclass(defn)
 
-        for decorator in defn.decorators:
-            self.analyze_class_decorator(defn, decorator)
+            for decorator in defn.decorators:
+                self.analyze_class_decorator(defn, decorator)
 
-        self.enter_class(defn)
+            self.enter_class(defn)
 
-        # Analyze class body.
-        defn.defs.accept(self)
+            # Analyze class body.
+            defn.defs.accept(self)
 
-        self.calculate_abstract_status(defn.info)
-        self.setup_type_promotion(defn)
+            self.calculate_abstract_status(defn.info)
+            self.setup_type_promotion(defn)
 
-        self.leave_class()
-        self.unbind_class_type_vars()
+            self.leave_class()
+
+            self.unbind_class_type_vars()
 
     def enter_class(self, defn: ClassDef) -> None:
         # Remember previous active class
@@ -818,21 +823,24 @@ def analyze_namedtuple_classdef(self, defn: ClassDef) -> bool:
                     node = self.lookup(defn.name, defn)
                     if node is not None:
                         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)
+                        items, types, default_items = self.check_namedtuple_classdef(defn)
+                        node.node = self.build_namedtuple_typeinfo(
+                            defn.name, items, types, default_items)
                         return True
         return False
 
-    def check_namedtuple_classdef(self, defn: ClassDef) -> Tuple[List[str], List[Type]]:
+    def check_namedtuple_classdef(
+            self, defn: ClassDef) -> Tuple[List[str], List[Type], Dict[str, Expression]]:
         NAMEDTUP_CLASS_ERROR = ('Invalid statement in NamedTuple definition; '
                                'expected "field_name: field_type"')
         if self.options.python_version < (3, 6):
             self.fail('NamedTuple class syntax is only supported in Python 3.6', defn)
-            return [], []
+            return [], [], {}
         if len(defn.base_type_exprs) > 1:
             self.fail('NamedTuple should be a single base', defn)
         items = []  # type: List[str]
         types = []  # type: List[Type]
+        default_items = {}  # type: Dict[str, Expression]
         for stmt in defn.defs.body:
             if not isinstance(stmt, AssignmentStmt):
                 # Still allow pass or ... (for empty namedtuples).
@@ -854,10 +862,14 @@ def check_namedtuple_classdef(self, defn: ClassDef) -> Tuple[List[str], List[Typ
                               .format(name), stmt)
                 if stmt.type is None or hasattr(stmt, 'new_syntax') and not stmt.new_syntax:
                     self.fail(NAMEDTUP_CLASS_ERROR, stmt)
-                elif not isinstance(stmt.rvalue, TempNode):
+                elif isinstance(stmt.rvalue, TempNode):
                     # x: int assigns rvalue to TempNode(AnyType())
-                    self.fail('Right hand side values are not supported in NamedTuple', stmt)
-        return items, types
+                    if default_items:
+                        self.fail('Non-default NamedTuple fields cannot follow default fields',
+                                  stmt)
+                else:
+                    default_items[name] = stmt.rvalue
+        return items, types, default_items
 
     def setup_class_def_analysis(self, defn: ClassDef) -> None:
         """Prepare for the analysis of a class definition."""
@@ -1915,12 +1927,12 @@ def check_namedtuple(self, node: Expression, var_name: str = None) -> Optional[T
         items, types, ok = self.parse_namedtuple_args(call, fullname)
         if not ok:
             # Error. Construct dummy return value.
-            return self.build_namedtuple_typeinfo('namedtuple', [], [])
+            return self.build_namedtuple_typeinfo('namedtuple', [], [], {})
         name = cast(StrExpr, call.args[0]).value
         if name != var_name or self.is_func_scope():
             # Give it a unique name derived from the line number.
             name += '@' + str(call.line)
-        info = self.build_namedtuple_typeinfo(name, items, types)
+        info = self.build_namedtuple_typeinfo(name, items, types, {})
         # Store it as a global just in case it would remain anonymous.
         # (Or in the nearest class if there is one.)
         stnode = SymbolTableNode(GDEF, info, self.cur_mod_id)
@@ -2013,8 +2025,8 @@ def basic_new_typeinfo(self, name: str, basetype_or_fallback: Instance) -> TypeI
         info.bases = [basetype_or_fallback]
         return info
 
-    def build_namedtuple_typeinfo(self, name: str, items: List[str],
-                                  types: List[Type]) -> TypeInfo:
+    def build_namedtuple_typeinfo(self, name: str, items: List[str], types: List[Type],
+                                  default_items: Dict[str, Expression]) -> TypeInfo:
         strtype = self.str_type()
         basetuple_type = self.named_type('__builtins__.tuple', [AnyType()])
         dictype = (self.named_type_or_none('builtins.dict', [strtype, AnyType()])
@@ -2046,6 +2058,7 @@ def add_field(var: Var, is_initialized_in_class: bool = False,
         tuple_of_strings = TupleType([strtype for _ in items], basetuple_type)
         add_field(Var('_fields', tuple_of_strings), is_initialized_in_class=True)
         add_field(Var('_field_types', dictype), is_initialized_in_class=True)
+        add_field(Var('_field_defaults', dictype), is_initialized_in_class=True)
         add_field(Var('_source', strtype), is_initialized_in_class=True)
 
         tvd = TypeVarDef('NT', 1, [], info.tuple_type)
@@ -2083,8 +2096,14 @@ def add_method(funcname: str,
 
         add_method('_replace', ret=selftype,
                    args=[Argument(var, var.type, EllipsisExpr(), ARG_NAMED_OPT) for var in vars])
+
+        def make_init_arg(var: Var) -> Argument:
+            default = default_items.get(var.name(), None)
+            kind = ARG_POS if default is None else ARG_OPT
+            return Argument(var, var.type, default, kind)
+
         add_method('__init__', ret=NoneTyp(), name=info.name(),
-                   args=[Argument(var, var.type, None, ARG_POS) for var in vars])
+                   args=[make_init_arg(var) for var in vars])
         add_method('_asdict', args=[], ret=ordereddictype)
         add_method('_make', ret=selftype, is_classmethod=True,
                    args=[Argument(Var('iterable', iterable_type), iterable_type, None, ARG_POS),

test-data/unit/check-class-namedtuple.test

@@ -294,6 +294,10 @@ class X(NamedTuple):
     y: str
 
 reveal_type(X._fields)  # E: Revealed type is 'Tuple[builtins.str, builtins.str]'
+reveal_type(X._field_types)  # E: Revealed type is 'builtins.dict[builtins.str, Any]'
+reveal_type(X._field_defaults)  # E: Revealed type is 'builtins.dict[builtins.str, Any]'
+
+[builtins fixtures/dict.pyi]
 
 [case testNewNamedTupleUnit]
 # flags: --python-version 3.6
@@ -349,9 +353,17 @@ import typing
 
 class X(typing.NamedTuple):
     x: int
-    y: str = 'y'  # E: Right hand side values are not supported in NamedTuple
-    z = None  # type: int # E: Invalid statement in NamedTuple definition; expected "field_name: field_type"
-    x[0]: int  # E: Invalid statement in NamedTuple definition; expected "field_name: field_type"
+    y = 1
+    x.x: int
+    z: str = 'z'
+    aa: int
+
+[out]
+main:6: error: Invalid statement in NamedTuple definition; expected "field_name: field_type"
+main:7: error: Invalid statement in NamedTuple definition; expected "field_name: field_type"
+main:7: error: Type cannot be declared in assignment to non-self attribute
+main:7: error: "int" has no attribute "x"
+main:9: error: Non-default NamedTuple fields cannot follow default fields
 
 [builtins fixtures/list.pyi]
 
@@ -376,3 +388,97 @@ def f(a: Type[N]):
 [builtins fixtures/list.pyi]
 [out]
 main:8: error: Unsupported type Type["N"]
+
+[case testNewNamedTupleWithDefaults]
+# flags: --fast-parser --python-version 3.6
+from typing import List, NamedTuple, Optional
+
+class X(NamedTuple):
+    x: int
+    y: int = 2
+
+reveal_type(X(1))  # E: Revealed type is 'Tuple[builtins.int, builtins.int, fallback=__main__.X]'
+reveal_type(X(1, 2))  # E: Revealed type is 'Tuple[builtins.int, builtins.int, fallback=__main__.X]'
+
+X(1, 'a')  # E: Argument 2 to "X" has incompatible type "str"; expected "int"
+X(1, z=3)  # E: Unexpected keyword argument "z" for "X"
+
+class HasNone(NamedTuple):
+    x: int
+    y: Optional[int] = None
+
+reveal_type(HasNone(1))  # E: Revealed type is 'Tuple[builtins.int, builtins.int, fallback=__main__.HasNone]'
+
+class Parameterized(NamedTuple):
+    x: int
+    y: List[int] = [1] + [2]
+    z: List[int] = []
+
+reveal_type(Parameterized(1))  # E: Revealed type is 'Tuple[builtins.int, builtins.list[builtins.int], builtins.list[builtins.int], fallback=__main__.Parameterized]'
+Parameterized(1, ['not an int'])  # E: List item 0 has incompatible type "str"
+
+class Default:
+    pass
+
+class UserDefined(NamedTuple):
+    x: Default = Default()
+
+reveal_type(UserDefined())  # E: Revealed type is 'Tuple[__main__.Default, fallback=__main__.UserDefined]'
+reveal_type(UserDefined(Default()))  # E: Revealed type is 'Tuple[__main__.Default, fallback=__main__.UserDefined]'
+UserDefined(1)  # E: Argument 1 to "UserDefined" has incompatible type "int"; expected "Default"
+
+[builtins fixtures/list.pyi]
+
+[case testNewNamedTupleWithDefaultsStrictOptional]
+# flags: --fast-parser --strict-optional --python-version 3.6
+from typing import List, NamedTuple, Optional
+
+class HasNone(NamedTuple):
+    x: int
+    y: Optional[int] = None
+
+reveal_type(HasNone(1))  # E: Revealed type is 'Tuple[builtins.int, Union[builtins.int, builtins.None], fallback=__main__.HasNone]'
+HasNone(None)  # E: Argument 1 to "HasNone" has incompatible type None; expected "int"
+HasNone(1, y=None)
+HasNone(1, y=2)
+
+class CannotBeNone(NamedTuple):
+    x: int
+    y: int = None  # E: Incompatible types in assignment (expression has type None, variable has type "int")
+
+[builtins fixtures/list.pyi]
+
+[case testNewNamedTupleWrongType]
+# flags: --fast-parser --python-version 3.6
+from typing import NamedTuple
+
+class X(NamedTuple):
+    x: int
+    y: int = 'not an int'  # E: Incompatible types in assignment (expression has type "str", variable has type "int")
+
+[case testNewNamedTupleErrorInDefault]
+# flags: --fast-parser --python-version 3.6
+from typing import NamedTuple
+
+class X(NamedTuple):
+    x: int = 1 + '1'  # E: Unsupported operand types for + ("int" and "str")
+
+[case testNewNamedTupleInheritance]
+# flags: --fast-parser --python-version 3.6
+from typing import NamedTuple
+
+class X(NamedTuple):
+    x: str
+    y: int = 3
+
+class Y(X):
+    def method(self) -> str:
+        self.y
+        return self.x
+
+reveal_type(Y('a'))  # E: Revealed type is 'Tuple[builtins.str, builtins.int, fallback=__main__.Y]'
+Y(y=1, x='1').method()
+
+class CallsBaseInit(X):
+    def __init__(self, x: str) -> None:
+        super().__init__(x)