Add support for NamedTuple methods

mypy/semanal.py

@@ -156,6 +156,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.
@@ -304,7 +311,7 @@ def file_context(self, file_node: MypyFile, fnam: str, options: Options,
         self.is_stub_file = fnam.lower().endswith('.pyi')
         self.globals = file_node.names
         if active_type:
-            self.enter_class(active_type.defn)
+            self.enter_class(active_type.defn.info)
             # TODO: Bind class type vars
 
         yield
@@ -626,11 +633,32 @@ def analyze_class_body(self, defn: ClassDef) -> Iterator[bool]:
             if self.analyze_typeddict_classdef(defn):
                 yield False
                 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.enter_class(named_tuple_info)
+
                 yield True
+
                 self.leave_class()
+
+                # 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.analyze_base_classes(defn)
@@ -639,21 +667,21 @@ def analyze_class_body(self, defn: ClassDef) -> Iterator[bool]:
                 for decorator in defn.decorators:
                     self.analyze_class_decorator(defn, decorator)
 
-                self.enter_class(defn)
+                self.enter_class(defn.info)
                 yield True
 
                 self.calculate_abstract_status(defn.info)
                 self.setup_type_promotion(defn)
 
                 self.leave_class()
 
-    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. """
@@ -805,7 +833,7 @@ def get_all_bases_tvars(self, defn: ClassDef, removed: List[int]) -> TypeVarList
                 tvars.extend(base_tvars)
         return remove_dups(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):
@@ -815,15 +843,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 [= default]"')
         if self.options.python_version < (3, 6):
             self.fail('NamedTuple class syntax is only supported in Python 3.6', defn)
             return [], [], {}
@@ -835,10 +865,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)
@@ -2066,6 +2104,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)
@@ -3437,7 +3477,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)
@@ -3579,8 +3619,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]
 
@@ -344,8 +345,8 @@ 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"
+    y = z = 2  # E: Invalid statement in NamedTuple definition; expected "field_name: field_type [= default]"
+    def f(self): pass
 
 [case testNewNamedTupleWithInvalidItems2]
 # flags: --python-version 3.6
@@ -359,8 +360,8 @@ class X(typing.NamedTuple):
     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:6: error: Invalid statement in NamedTuple definition; expected "field_name: field_type [= default]"
+main:7: error: Invalid statement in NamedTuple definition; expected "field_name: field_type [= default]"
 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
@@ -373,7 +374,7 @@ from typing import NamedTuple
 
 class X(NamedTuple):
     x: int
-    y = 2  # E: Invalid statement in NamedTuple definition; expected "field_name: field_type"
+    y = 2  # E: Invalid statement in NamedTuple definition; expected "field_name: field_type [= default]"
 
 [case testTypeUsingTypeCNamedTuple]
 # 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 [= default]" \
+        # 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'