Support additional args to namedtuple() (#5215)

Author: JelleZijlstra

mypy/checkexpr.py

@@ -231,8 +231,15 @@ def analyze_var_ref(self, var: Var, context: Context) -> Type:
     def visit_call_expr(self, e: CallExpr, allow_none_return: bool = False) -> Type:
         """Type check a call expression."""
         if e.analyzed:
+            if isinstance(e.analyzed, NamedTupleExpr) and not e.analyzed.is_typed:
+                # Type check the arguments, but ignore the results. This relies
+                # on the typeshed stubs to type check the arguments.
+                self.visit_call_expr_inner(e)
             # It's really a special form that only looks like a call.
             return self.accept(e.analyzed, self.type_context[-1])
+        return self.visit_call_expr_inner(e, allow_none_return=allow_none_return)
+
+    def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> Type:
         if isinstance(e.callee, NameExpr) and isinstance(e.callee.node, TypeInfo) and \
                 e.callee.node.typeddict_type is not None:
             # Use named fallback for better error messages.

mypy/nodes.py

@@ -1917,10 +1917,12 @@ class NamedTupleExpr(Expression):
     # The class representation of this named tuple (its tuple_type attribute contains
     # the tuple item types)
     info = None  # type: TypeInfo
+    is_typed = False  # whether this class was created with typing.NamedTuple
 
-    def __init__(self, info: 'TypeInfo') -> None:
+    def __init__(self, info: 'TypeInfo', is_typed: bool = False) -> None:
         super().__init__()
         self.info = info
+        self.is_typed = is_typed
 
     def accept(self, visitor: ExpressionVisitor[T]) -> T:
         return visitor.visit_namedtuple_expr(self)

mypy/semanal_namedtuple.py

@@ -3,7 +3,7 @@
 This is conceptually part of mypy.semanal (semantic analyzer pass 2).
 """
 
-from typing import Tuple, List, Dict, Optional, cast
+from typing import Tuple, List, Dict, Mapping, Optional, cast
 
 from mypy.types import (
     Type, TupleType, NoneTyp, AnyType, TypeOfAny, TypeVarType, TypeVarDef, CallableType, TypeType
@@ -14,7 +14,7 @@
     AssignmentStmt, PassStmt, Decorator, FuncBase, ClassDef, Expression, RefExpr, TypeInfo,
     NamedTupleExpr, CallExpr, Context, TupleExpr, ListExpr, SymbolTableNode, FuncDef, Block,
     TempNode,
-    ARG_POS, ARG_NAMED_OPT, ARG_OPT, MDEF, GDEF
+    ARG_POS, ARG_NAMED, ARG_NAMED_OPT, ARG_OPT, MDEF, GDEF
 )
 from mypy.options import Options
 from mypy.exprtotype import expr_to_unanalyzed_type, TypeTranslationError
@@ -48,7 +48,7 @@ def analyze_namedtuple_classdef(self, defn: ClassDef) -> Optional[TypeInfo]:
                         node.node = info
                         defn.info.replaced = info
                         defn.info = info
-                        defn.analyzed = NamedTupleExpr(info)
+                        defn.analyzed = NamedTupleExpr(info, is_typed=True)
                         defn.analyzed.line = defn.line
                         defn.analyzed.column = defn.column
                         return info
@@ -142,35 +142,71 @@ def check_namedtuple(self,
         if not isinstance(callee, RefExpr):
             return None
         fullname = callee.fullname
-        if fullname not in ('collections.namedtuple', 'typing.NamedTuple'):
+        if fullname == 'collections.namedtuple':
+            is_typed = False
+        elif fullname == 'typing.NamedTuple':
+            is_typed = True
+        else:
             return None
-        items, types, ok = self.parse_namedtuple_args(call, fullname)
+        items, types, defaults, ok = self.parse_namedtuple_args(call, fullname)
         if not ok:
             # Error. Construct dummy return value.
             return self.build_namedtuple_typeinfo('namedtuple', [], [], {})
         name = cast(StrExpr, call.args[0]).value
         if name != var_name or 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, {})
+        if len(defaults) > 0:
+            default_items = {
+                arg_name: default
+                for arg_name, default in zip(items[-len(defaults):], defaults)
+            }
+        else:
+            default_items = {}
+        info = self.build_namedtuple_typeinfo(name, items, types, default_items)
         # 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.api.add_symbol_table_node(name, stnode)
-        call.analyzed = NamedTupleExpr(info)
+        call.analyzed = NamedTupleExpr(info, is_typed=is_typed)
         call.analyzed.set_line(call.line, call.column)
         return info
 
-    def parse_namedtuple_args(self, call: CallExpr,
-                              fullname: str) -> Tuple[List[str], List[Type], bool]:
+    def parse_namedtuple_args(self, call: CallExpr, fullname: str
+                              ) -> Tuple[List[str], List[Type], List[Expression], bool]:
+        """Parse a namedtuple() call into data needed to construct a type.
+
+        Returns a 4-tuple:
+        - List of argument names
+        - List of argument types
+        - Number of arguments that have a default value
+        - Whether the definition typechecked.
+
+        """
         # TODO: Share code with check_argument_count in checkexpr.py?
         args = call.args
         if len(args) < 2:
             return self.fail_namedtuple_arg("Too few arguments for namedtuple()", call)
+        defaults = []  # type: List[Expression]
         if len(args) > 2:
-            # FIX incorrect. There are two additional parameters
-            return self.fail_namedtuple_arg("Too many arguments for namedtuple()", call)
-        if call.arg_kinds != [ARG_POS, ARG_POS]:
+            # Typed namedtuple doesn't support additional arguments.
+            if fullname == 'typing.NamedTuple':
+                return self.fail_namedtuple_arg("Too many arguments for NamedTuple()", call)
+            for i, arg_name in enumerate(call.arg_names[2:], 2):
+                if arg_name == 'defaults':
+                    arg = args[i]
+                    # We don't care what the values are, as long as the argument is an iterable
+                    # and we can count how many defaults there are.
+                    if isinstance(arg, (ListExpr, TupleExpr)):
+                        defaults = list(arg.items)
+                    else:
+                        self.fail(
+                            "List or tuple literal expected as the defaults argument to "
+                            "namedtuple()",
+                            arg
+                        )
+                    break
+        if call.arg_kinds[:2] != [ARG_POS, ARG_POS]:
             return self.fail_namedtuple_arg("Unexpected arguments to namedtuple()", call)
         if not isinstance(args[0], (StrExpr, BytesExpr, UnicodeExpr)):
             return self.fail_namedtuple_arg(
@@ -196,17 +232,21 @@ def parse_namedtuple_args(self, call: CallExpr,
                 items = [cast(StrExpr, item).value for item in listexpr.items]
             else:
                 # The fields argument contains (name, type) tuples.
-                items, types, ok = self.parse_namedtuple_fields_with_types(listexpr.items, call)
+                items, types, _, ok = self.parse_namedtuple_fields_with_types(listexpr.items, call)
         if not types:
             types = [AnyType(TypeOfAny.unannotated) for _ in items]
         underscore = [item for item in items if item.startswith('_')]
         if underscore:
             self.fail("namedtuple() field names cannot start with an underscore: "
                       + ', '.join(underscore), call)
-        return items, types, ok
+        if len(defaults) > len(items):
+            self.fail("Too many defaults given in call to namedtuple()", call)
+            defaults = defaults[:len(items)]
+        return items, types, defaults, ok
 
-    def parse_namedtuple_fields_with_types(self, nodes: List[Expression],
-                                           context: Context) -> Tuple[List[str], List[Type], bool]:
+    def parse_namedtuple_fields_with_types(self, nodes: List[Expression], context: Context
+                                           ) -> Tuple[List[str], List[Type], List[Expression],
+                                                      bool]:
         items = []  # type: List[str]
         types = []  # type: List[Type]
         for item in nodes:
@@ -226,15 +266,15 @@ def parse_namedtuple_fields_with_types(self, nodes: List[Expression],
                 types.append(self.api.anal_type(type))
             else:
                 return self.fail_namedtuple_arg("Tuple expected as NamedTuple() field", item)
-        return items, types, True
+        return items, types, [], True
 
-    def fail_namedtuple_arg(self, message: str,
-                            context: Context) -> Tuple[List[str], List[Type], bool]:
+    def fail_namedtuple_arg(self, message: str, context: Context
+                            ) -> Tuple[List[str], List[Type], List[Expression], bool]:
         self.fail(message, context)
-        return [], [], False
+        return [], [], [], False
 
     def build_namedtuple_typeinfo(self, name: str, items: List[str], types: List[Type],
-                                  default_items: Dict[str, Expression]) -> TypeInfo:
+                                  default_items: Mapping[str, Expression]) -> TypeInfo:
         strtype = self.api.named_type('__builtins__.str')
         implicit_any = AnyType(TypeOfAny.special_form)
         basetuple_type = self.api.named_type('__builtins__.tuple', [implicit_any])

test-data/unit/check-namedtuple.test

@@ -1,7 +1,7 @@
 [case testNamedTupleUsedAsTuple]
 from collections import namedtuple
 
-X = namedtuple('X', ['x', 'y'])
+X = namedtuple('X', 'x y')
 x = None  # type: X
 a, b = x
 b = x[0]
@@ -28,7 +28,7 @@ X = namedtuple('X', 'x, _y, _z')  # E: namedtuple() field names cannot start wit
 [case testNamedTupleAccessingAttributes]
 from collections import namedtuple
 
-X = namedtuple('X', ['x', 'y'])
+X = namedtuple('X', 'x y')
 x = None  # type: X
 x.x
 x.y
@@ -38,7 +38,7 @@ x.z # E: "X" has no attribute "z"
 [case testNamedTupleAttributesAreReadOnly]
 from collections import namedtuple
 
-X = namedtuple('X', ['x', 'y'])
+X = namedtuple('X', 'x y')
 x = None  # type: 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
@@ -71,7 +71,7 @@ main:10: note: Protocol member HasX.x expected settable variable, got read-only
 [case testNamedTupleCreateWithPositionalArguments]
 from collections import namedtuple
 
-X = namedtuple('X', ['x', 'y'])
+X = namedtuple('X', 'x y')
 x = X(1, 'x')
 x.x
 x.z      # E: "X" has no attribute "z"
@@ -81,21 +81,52 @@ x = X(1, 2, 3)  # E: Too many arguments for "X"
 [case testCreateNamedTupleWithKeywordArguments]
 from collections import namedtuple
 
-X = namedtuple('X', ['x', 'y'])
+X = namedtuple('X', 'x y')
 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"
 
-
 [case testNamedTupleCreateAndUseAsTuple]
 from collections import namedtuple
 
-X = namedtuple('X', ['x', 'y'])
+X = namedtuple('X', 'x y')
 x = X(1, 'x')
 a, b = x
 a, b, c = x  # E: Need more than 2 values to unpack (3 expected)
 
+[case testNamedTupleAdditionalArgs]
+from collections import namedtuple
+
+A = namedtuple('A', 'a b')
+B = namedtuple('B', 'a b', rename=1)
+C = namedtuple('C', 'a b', rename='not a bool')
+D = namedtuple('D', 'a b', unrecognized_arg=False)
+E = namedtuple('E', 'a b', 0)
+
+[builtins fixtures/bool.pyi]
+
+[out]
+main:5: error: Argument "rename" to "namedtuple" has incompatible type "str"; expected "int"
+main:6: error: Unexpected keyword argument "unrecognized_arg" for "namedtuple"
+<ROOT>/test-data/unit/lib-stub/collections.pyi:3: note: "namedtuple" defined here
+main:7: error: Too many positional arguments for "namedtuple"
+
+[case testNamedTupleDefaults]
+# flags: --python-version 3.7
+from collections import namedtuple
+
+X = namedtuple('X', ['x', 'y'], defaults=(1,))
+
+X()  # E: Too few arguments for "X"
+X(0)  # ok
+X(0, 1)  # ok
+X(0, 1, 2)  # E: Too many arguments for "X"
+
+Y = namedtuple('Y', ['x', 'y'], defaults=(1, 2, 3))  # E: Too many defaults given in call to namedtuple()
+Z = namedtuple('Z', ['x', 'y'], defaults='not a tuple')  # E: Argument "defaults" to "namedtuple" has incompatible type "str"; expected "Optional[Iterable[Any]]"  # E: List or tuple literal expected as the defaults argument to namedtuple()
+
+[builtins fixtures/list.pyi]
 
 [case testNamedTupleWithItemTypes]
 from typing import NamedTuple
@@ -240,6 +271,7 @@ import collections
 MyNamedTuple = collections.namedtuple('MyNamedTuple', ['spam', 'eggs'])
 MyNamedTuple.x # E: "Type[MyNamedTuple]" has no attribute "x"
 
+[builtins fixtures/list.pyi]
 
 [case testNamedTupleEmptyItems]
 from typing import NamedTuple
@@ -280,6 +312,8 @@ x._replace(x=3, y=5)
 x._replace(z=5)  # E: Unexpected keyword argument "z" for "_replace" of "X"
 x._replace(5)  # E: Too many positional arguments for "_replace" of "X"
 
+[builtins fixtures/list.pyi]
+
 [case testNamedTupleReplaceAsClass]
 from collections import namedtuple
 
@@ -288,6 +322,7 @@ x = None  # type: X
 X._replace(x, x=1, y=2)
 X._replace(x=1, y=2)  # E: Missing positional argument "self" in call to "_replace" of "X"
 
+[builtins fixtures/list.pyi]
 
 [case testNamedTupleReplaceTyped]
 from typing import NamedTuple

test-data/unit/check-newtype.test

@@ -123,6 +123,9 @@ Point3 = NewType('Point3', Vector3)
 p3 = Point3(Vector3(1, 3))
 reveal_type(p3.x)  # E: Revealed type is 'builtins.int'
 reveal_type(p3.y)  # E: Revealed type is 'builtins.int'
+
+[builtins fixtures/list.pyi]
+
 [out]
 
 [case testNewTypeWithCasts]

test-data/unit/check-python2.test

@@ -13,14 +13,6 @@ s = b'foo'
 from typing import TypeVar
 T = TypeVar(u'T')
 
-[case testNamedTupleUnicode]
-from typing import NamedTuple
-from collections import namedtuple
-N = NamedTuple(u'N', [(u'x', int)])
-n = namedtuple(u'n', u'x y')
-
-[builtins fixtures/dict.pyi]
-
 [case testPrintStatement]
 print ''() # E: "str" not callable
 print 1, 1() # E: "int" not callable

test-data/unit/fixtures/args.pyi

@@ -27,3 +27,4 @@ class int:
 class str: pass
 class bool: pass
 class function: pass
+class ellipsis: pass

test-data/unit/fixtures/ops.pyi

@@ -55,3 +55,5 @@ class float: pass
 class BaseException: pass
 
 def __print(a1=None, a2=None, a3=None, a4=None): pass
+
+class ellipsis: pass

test-data/unit/lib-stub/collections.pyi

@@ -1,3 +1,11 @@
-import typing
+from typing import Any, Iterable, Union, Optional
 
-namedtuple = object()
+def namedtuple(
+    typename: str,
+    field_names: Union[str, Iterable[str]],
+    *,
+    # really bool but many tests don't have bool available
+    rename: int = ...,
+    module: Optional[str] = ...,
+    defaults: Optional[Iterable[Any]] = ...
+) -> Any: ...

test-data/unit/python2eval.test

@@ -262,13 +262,17 @@ s = ''.join([u''])  # Error
 _program.py:5: error: Incompatible types in assignment (expression has type "str", variable has type "int")
 _program.py:6: error: Incompatible types in assignment (expression has type "unicode", variable has type "str")
 
-[case testNamedTupleError_python2]
-import typing
+[case testNamedTuple_python2]
+from typing import NamedTuple
 from collections import namedtuple
 X = namedtuple('X', ['a', 'b'])
 x = X(a=1, b='s')
 x.c
 x.a
+
+N = NamedTuple(u'N', [(u'x', int)])
+n = namedtuple(u'n', u'x y')
+
 [out]
 _program.py:5: error: "X" has no attribute "c"