Add support for NamedTuple methods

mypy/semanal.py

@@ -151,6 +151,13 @@
 FUNCTION_FIRST_PHASE_POSTPONE_SECOND = 1  # Add to symbol table but postpone body
 FUNCTION_SECOND_PHASE = 2  # Only analyze body
 
+# Matches "_prohibited" in typing.py, but adds __annotations__, which works at runtime but can't
+# easily be supported in a static checker.
+NAMEDTUPLE_PROHIBITED_NAMES = ('__new__', '__init__', '__slots__', '__getnewargs__',
+                               '_fields', '_field_defaults', '_field_types',
+                               '_make', '_replace', '_asdict', '_source',
+                               '__annotations__')
+
 
 class SemanticAnalyzer(NodeVisitor):
     """Semantically analyze parsed mypy files.
@@ -655,23 +662,46 @@ def visit_class_def(self, defn: ClassDef) -> None:
         self.clean_up_bases_and_infer_type_variables(defn)
         if self.analyze_typeddict_classdef(defn):
             return
-        if self.analyze_namedtuple_classdef(defn):
-            # just analyze the class body so we catch type errors in default values
-            self.enter_class(defn)
+        named_tuple_info = self.analyze_namedtuple_classdef(defn)
+        if named_tuple_info is not None:
+            # Temporarily clear the names dict so we don't get errors about duplicate names that
+            # were already set in build_namedtuple_typeinfo.
+            nt_names = named_tuple_info.names
+            named_tuple_info.names = SymbolTable()
+
+            self.bind_class_type_vars(named_tuple_info)
+            self.enter_class(named_tuple_info)
+
             defn.defs.accept(self)
+
             self.leave_class()
+            self.unbind_class_type_vars()
+
+            # make sure we didn't use illegal names, then reset the names in the typeinfo
+            for prohibited in NAMEDTUPLE_PROHIBITED_NAMES:
+                if prohibited in named_tuple_info.names:
+                    self.fail('Cannot overwrite NamedTuple attribute "{}"'.format(prohibited),
+                              named_tuple_info.names[prohibited].node)
+
+            # Restore the names in the original symbol table. This ensures that the symbol
+            # table contains the field objects created by build_namedtuple_typeinfo. Exclude
+            # __doc__, which can legally be overwritten by the class.
+            named_tuple_info.names.update({
+                key: value for key, value in nt_names.items()
+                if key not in named_tuple_info.names or key != '__doc__'
+            })
         else:
             self.setup_class_def_analysis(defn)
 
-            self.bind_class_type_vars(defn)
+            self.bind_class_type_vars(defn.info)
 
             self.analyze_base_classes(defn)
             self.analyze_metaclass(defn)
 
             for decorator in defn.decorators:
                 self.analyze_class_decorator(defn, decorator)
 
-            self.enter_class(defn)
+            self.enter_class(defn.info)
 
             # Analyze class body.
             defn.defs.accept(self)
@@ -683,13 +713,13 @@ def visit_class_def(self, defn: ClassDef) -> None:
 
             self.unbind_class_type_vars()
 
-    def enter_class(self, defn: ClassDef) -> None:
+    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.block_depth.append(-1)  # The class body increments this to 0
         self.postpone_nested_functions_stack.append(FUNCTION_BOTH_PHASES)
-        self.type = defn.info
+        self.type = info
 
     def leave_class(self) -> None:
         """ Restore analyzer state. """
@@ -698,14 +728,14 @@ def leave_class(self) -> None:
         self.locals.pop()
         self.type = self.type_stack.pop()
 
-    def bind_class_type_vars(self, defn: ClassDef) -> None:
+    def bind_class_type_vars(self, info: TypeInfo) -> None:
         """ Unbind type variables of previously active class and bind
         the type variables for the active class.
         """
         if self.bound_tvars:
             disable_typevars(self.bound_tvars)
         self.tvar_stack.append(self.bound_tvars)
-        self.bound_tvars = self.bind_class_type_variables_in_symbol_table(defn.info)
+        self.bound_tvars = self.bind_class_type_variables_in_symbol_table(info)
 
     def unbind_class_type_vars(self) -> None:
         """ Unbind the active class' type vars and rebind the
@@ -882,7 +912,7 @@ def remove_dups(self, tvars: List[T]) -> List[T]:
                 all_tvars.remove(t)
         return new_tvars
 
-    def analyze_namedtuple_classdef(self, defn: ClassDef) -> bool:
+    def analyze_namedtuple_classdef(self, defn: ClassDef) -> Optional[TypeInfo]:
         # special case for NamedTuple
         for base_expr in defn.base_type_exprs:
             if isinstance(base_expr, RefExpr):
@@ -892,15 +922,17 @@ def analyze_namedtuple_classdef(self, defn: ClassDef) -> bool:
                     if node is not None:
                         node.kind = GDEF  # TODO in process_namedtuple_definition also applies here
                         items, types, default_items = self.check_namedtuple_classdef(defn)
-                        node.node = self.build_namedtuple_typeinfo(
+                        info = self.build_namedtuple_typeinfo(
                             defn.name, items, types, default_items)
-                        return True
-        return False
+                        node.node = info
+                        defn.info = info
+                        return info
+        return None
 
     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"')
+                                '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 [], [], {}
@@ -912,10 +944,18 @@ def check_namedtuple_classdef(
         for stmt in defn.defs.body:
             if not isinstance(stmt, AssignmentStmt):
                 # Still allow pass or ... (for empty namedtuples).
-                if (not isinstance(stmt, PassStmt) and
-                    not (isinstance(stmt, ExpressionStmt) and
-                         isinstance(stmt.expr, EllipsisExpr))):
-                    self.fail(NAMEDTUP_CLASS_ERROR, stmt)
+                if (isinstance(stmt, PassStmt) or
+                    (isinstance(stmt, ExpressionStmt) and
+                        isinstance(stmt.expr, EllipsisExpr))):
+                    continue
+                # Also allow methods.
+                if isinstance(stmt, FuncBase):
+                    continue
+                # And docstrings.
+                if (isinstance(stmt, ExpressionStmt) and
+                        isinstance(stmt.expr, StrExpr)):
+                    continue
+                self.fail(NAMEDTUP_CLASS_ERROR, stmt)
             elif len(stmt.lvalues) > 1 or not isinstance(stmt.lvalues[0], NameExpr):
                 # An assignment, but an invalid one.
                 self.fail(NAMEDTUP_CLASS_ERROR, stmt)
@@ -2128,6 +2168,8 @@ def add_field(var: Var, is_initialized_in_class: bool = False,
         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)
+        add_field(Var('__annotations__', ordereddictype), is_initialized_in_class=True)
+        add_field(Var('__doc__', strtype), is_initialized_in_class=True)
 
         tvd = TypeVarDef('NT', 1, [], info.tuple_type)
         selftype = TypeVarType(tvd)
@@ -3359,7 +3401,7 @@ def visit_class_def(self, cdef: ClassDef) -> None:
         self.process_nested_classes(cdef)
 
     def process_nested_classes(self, outer_def: ClassDef) -> None:
-        self.sem.enter_class(outer_def)
+        self.sem.enter_class(outer_def.info)
         for node in outer_def.defs.body:
             if isinstance(node, ClassDef):
                 node.info = TypeInfo(SymbolTable(), node, self.sem.cur_mod_id)
@@ -3488,8 +3530,11 @@ def visit_func_def(self, fdef: FuncDef) -> None:
         self.errors.pop_function()
 
     def visit_class_def(self, tdef: ClassDef) -> None:
-        for type in tdef.info.bases:
-            self.analyze(type)
+        # NamedTuple base classes are validated in check_namedtuple_classdef; we don't have to
+        # check them again here.
+        if not tdef.info.is_named_tuple:
+            for type in tdef.info.bases:
+                self.analyze(type)
         # Recompute MRO now that we have analyzed all modules, to pick
         # up superclasses of bases imported from other modules in an
         # import loop. (Only do so if we succeeded the first time.)

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

@@ -296,6 +296,7 @@ class X(NamedTuple):
 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]'
+reveal_type(X.__annotations__)  # E: Revealed type is 'builtins.dict[builtins.str, Any]'
 
 [builtins fixtures/dict.pyi]
 
@@ -345,7 +346,7 @@ from typing import NamedTuple
 class X(NamedTuple):
     x: int
     y = z = 2  # E: Invalid statement in NamedTuple definition; expected "field_name: field_type"
-    def f(self): pass  # E: Invalid statement in NamedTuple definition; expected "field_name: field_type"
+    def f(self): pass
 
 [case testNewNamedTupleWithInvalidItems2]
 # flags: --python-version 3.6
@@ -482,3 +483,146 @@ Y(y=1, x='1').method()
 class CallsBaseInit(X):
     def __init__(self, x: str) -> None:
         super().__init__(x)
+
+[case testNewNamedTupleWithMethods]
+from typing import NamedTuple
+
+class XMeth(NamedTuple):
+    x: int
+    def double(self) -> int:
+        return self.x
+    async def asyncdouble(self) -> int:
+        return self.x
+
+class XRepr(NamedTuple):
+    x: int
+    y: int = 1
+    def __str__(self) -> str:
+        return 'string'
+    def __add__(self, other: XRepr) -> int:
+        return 0
+
+reveal_type(XMeth(1).double()) # E: Revealed type is 'builtins.int'
+reveal_type(XMeth(1).asyncdouble())  # E: Revealed type is 'typing.Awaitable[builtins.int]'
+reveal_type(XMeth(42).x)  # E: Revealed type is 'builtins.int'
+reveal_type(XRepr(42).__str__())  # E: Revealed type is 'builtins.str'
+reveal_type(XRepr(1, 2).__add__(XRepr(3)))  # E: Revealed type is 'builtins.int'
+
+[case testNewNamedTupleOverloading]
+from typing import NamedTuple, overload
+
+class Overloader(NamedTuple):
+    x: int
+    @overload
+    def method(self, y: str) -> str: pass
+    @overload
+    def method(self, y: int) -> int: pass
+    def method(self, y):
+        return y
+
+reveal_type(Overloader(1).method('string'))  # E: Revealed type is 'builtins.str'
+reveal_type(Overloader(1).method(1))  # E: Revealed type is 'builtins.int'
+Overloader(1).method(('tuple',))  # E: No overload variant of "method" of "Overloader" matches argument types [Tuple[builtins.str]]
+
+[case testNewNamedTupleMethodInheritance]
+from typing import NamedTuple, TypeVar
+
+T = TypeVar('T')
+
+class Base(NamedTuple):
+    x: int
+    def copy(self: T) -> T:
+        reveal_type(self)  # E: Revealed type is 'T`-1'
+        return self
+    def good_override(self) -> int:
+        reveal_type(self)  # E: Revealed type is 'Tuple[builtins.int, fallback=__main__.Base]'
+        reveal_type(self[0])  # E: Revealed type is 'builtins.int'
+        self[0] = 3  # E: Unsupported target for indexed assignment
+        reveal_type(self.x)  # E: Revealed type is 'builtins.int'
+        self.x = 3  # E: Property "x" defined in "Base" is read-only
+        self[1]  # E: Tuple index out of range
+        self[T]  # E: Tuple index must be an integer literal
+        return self.x
+    def bad_override(self) -> int:
+        return self.x
+
+class Child(Base):
+    def new_method(self) -> int:
+        reveal_type(self)  # E: Revealed type is 'Tuple[builtins.int, fallback=__main__.Child]'
+        reveal_type(self[0])  # E: Revealed type is 'builtins.int'
+        self[0] = 3  # E: Unsupported target for indexed assignment
+        reveal_type(self.x)  # E: Revealed type is 'builtins.int'
+        self.x = 3  # E: Property "x" defined in "Child" is read-only
+        self[1]  # E: Tuple index out of range
+        return self.x
+    def good_override(self) -> int:
+        return 0
+    def bad_override(self) -> str:  # E: Return type of "bad_override" incompatible with supertype "Base"
+        return 'incompatible'
+
+def takes_base(base: Base) -> int:
+    return base.x
+
+reveal_type(Base(1).copy())  # E: Revealed type is 'Tuple[builtins.int, fallback=__main__.Base]'
+reveal_type(Child(1).copy())  # E: Revealed type is 'Tuple[builtins.int, fallback=__main__.Child]'
+reveal_type(Base(1).good_override())  # E: Revealed type is 'builtins.int'
+reveal_type(Child(1).good_override())  # E: Revealed type is 'builtins.int'
+reveal_type(Base(1).bad_override())  # E: Revealed type is 'builtins.int'
+reveal_type(takes_base(Base(1)))  # E: Revealed type is 'builtins.int'
+reveal_type(takes_base(Child(1)))  # E: Revealed type is 'builtins.int'
+
+[case testNewNamedTupleIllegalNames]
+from typing import Callable, NamedTuple
+
+class XMethBad(NamedTuple):
+    x: int
+    def _fields(self):  # E: Cannot overwrite NamedTuple attribute "_fields"
+        return 'no chance for this'
+
+class MagicalFields(NamedTuple):
+    x: int
+    def __slots__(self) -> None: pass  # E: Cannot overwrite NamedTuple attribute "__slots__"
+    def __new__(cls) -> None: pass  # E: Cannot overwrite NamedTuple attribute "__new__"
+    def _source(self) -> int: pass  # E: Cannot overwrite NamedTuple attribute "_source"
+    __annotations__ = {'x': float}  # E: NamedTuple field name cannot start with an underscore: __annotations__ \
+        # E: Invalid statement in NamedTuple definition; expected "field_name: field_type" \
+        # E: Cannot overwrite NamedTuple attribute "__annotations__"
+
+class AnnotationsAsAMethod(NamedTuple):
+    x: int
+    # This fails at runtime because typing.py assumes that __annotations__ is a dictionary.
+    def __annotations__(self) -> float:  # E: Cannot overwrite NamedTuple attribute "__annotations__"
+        return 1.0
+
+class ReuseNames(NamedTuple):
+    x: int
+    def x(self) -> str:  # E: Name 'x' already defined
+        return ''
+
+    def y(self) -> int:
+        return 0
+    y: str  # E: Name 'y' already defined
+
+class ReuseCallableNamed(NamedTuple):
+    z: Callable[[ReuseNames], int]
+    def z(self) -> int:  # E: Name 'z' already defined
+        return 0
+
+[builtins fixtures/dict.pyi]
+
+[case testNewNamedTupleDocString]
+from typing import NamedTuple
+
+class Documented(NamedTuple):
+    """This is a docstring."""
+    x: int
+
+reveal_type(Documented.__doc__)  # E: Revealed type is 'builtins.str'
+reveal_type(Documented(1).x)  # E: Revealed type is 'builtins.int'
+
+class BadDoc(NamedTuple):
+    x: int
+    def __doc__(self) -> str:
+        return ''
+
+reveal_type(BadDoc(1).__doc__())  # E: Revealed type is 'builtins.str'