Skip to content

Commit 3589c08

Browse files
JelleZijlstragvanrossum
authored andcommitted
Add support for NamedTuple methods (#3081)
This is the mypy implementation of python/typing#352. Also fixes #3075.
1 parent 6f9a355 commit 3589c08

File tree

3 files changed

+248
-33
lines changed

3 files changed

+248
-33
lines changed

mypy/semanal.py

Lines changed: 66 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,13 @@
156156
FUNCTION_FIRST_PHASE_POSTPONE_SECOND = 1 # Add to symbol table but postpone body
157157
FUNCTION_SECOND_PHASE = 2 # Only analyze body
158158

159+
# Matches "_prohibited" in typing.py, but adds __annotations__, which works at runtime but can't
160+
# easily be supported in a static checker.
161+
NAMEDTUPLE_PROHIBITED_NAMES = ('__new__', '__init__', '__slots__', '__getnewargs__',
162+
'_fields', '_field_defaults', '_field_types',
163+
'_make', '_replace', '_asdict', '_source',
164+
'__annotations__')
165+
159166

160167
class SemanticAnalyzer(NodeVisitor):
161168
"""Semantically analyze parsed mypy files.
@@ -304,7 +311,7 @@ def file_context(self, file_node: MypyFile, fnam: str, options: Options,
304311
self.is_stub_file = fnam.lower().endswith('.pyi')
305312
self.globals = file_node.names
306313
if active_type:
307-
self.enter_class(active_type.defn)
314+
self.enter_class(active_type.defn.info)
308315
# TODO: Bind class type vars
309316

310317
yield
@@ -631,11 +638,36 @@ def analyze_class_body(self, defn: ClassDef) -> Iterator[bool]:
631638
if self.analyze_typeddict_classdef(defn):
632639
yield False
633640
return
634-
if self.analyze_namedtuple_classdef(defn):
635-
# just analyze the class body so we catch type errors in default values
636-
self.enter_class(defn)
641+
named_tuple_info = self.analyze_namedtuple_classdef(defn)
642+
if named_tuple_info is not None:
643+
# Temporarily clear the names dict so we don't get errors about duplicate names
644+
# that were already set in build_namedtuple_typeinfo.
645+
nt_names = named_tuple_info.names
646+
named_tuple_info.names = SymbolTable()
647+
# This is needed for the cls argument to classmethods to get bound correctly.
648+
named_tuple_info.names['__init__'] = nt_names['__init__']
649+
650+
self.enter_class(named_tuple_info)
651+
637652
yield True
653+
638654
self.leave_class()
655+
656+
# make sure we didn't use illegal names, then reset the names in the typeinfo
657+
for prohibited in NAMEDTUPLE_PROHIBITED_NAMES:
658+
if prohibited in named_tuple_info.names:
659+
if nt_names.get(prohibited) is named_tuple_info.names[prohibited]:
660+
continue
661+
self.fail('Cannot overwrite NamedTuple attribute "{}"'.format(prohibited),
662+
named_tuple_info.names[prohibited].node)
663+
664+
# Restore the names in the original symbol table. This ensures that the symbol
665+
# table contains the field objects created by build_namedtuple_typeinfo. Exclude
666+
# __doc__, which can legally be overwritten by the class.
667+
named_tuple_info.names.update({
668+
key: value for key, value in nt_names.items()
669+
if key not in named_tuple_info.names or key != '__doc__'
670+
})
639671
else:
640672
self.setup_class_def_analysis(defn)
641673
self.analyze_base_classes(defn)
@@ -644,21 +676,21 @@ def analyze_class_body(self, defn: ClassDef) -> Iterator[bool]:
644676
for decorator in defn.decorators:
645677
self.analyze_class_decorator(defn, decorator)
646678

647-
self.enter_class(defn)
679+
self.enter_class(defn.info)
648680
yield True
649681

650682
self.calculate_abstract_status(defn.info)
651683
self.setup_type_promotion(defn)
652684

653685
self.leave_class()
654686

655-
def enter_class(self, defn: ClassDef) -> None:
687+
def enter_class(self, info: TypeInfo) -> None:
656688
# Remember previous active class
657689
self.type_stack.append(self.type)
658690
self.locals.append(None) # Add class scope
659691
self.block_depth.append(-1) # The class body increments this to 0
660692
self.postpone_nested_functions_stack.append(FUNCTION_BOTH_PHASES)
661-
self.type = defn.info
693+
self.type = info
662694

663695
def leave_class(self) -> None:
664696
""" Restore analyzer state. """
@@ -810,7 +842,7 @@ def get_all_bases_tvars(self, defn: ClassDef, removed: List[int]) -> TypeVarList
810842
tvars.extend(base_tvars)
811843
return remove_dups(tvars)
812844

813-
def analyze_namedtuple_classdef(self, defn: ClassDef) -> bool:
845+
def analyze_namedtuple_classdef(self, defn: ClassDef) -> Optional[TypeInfo]:
814846
# special case for NamedTuple
815847
for base_expr in defn.base_type_exprs:
816848
if isinstance(base_expr, RefExpr):
@@ -820,21 +852,17 @@ def analyze_namedtuple_classdef(self, defn: ClassDef) -> bool:
820852
if node is not None:
821853
node.kind = GDEF # TODO in process_namedtuple_definition also applies here
822854
items, types, default_items = self.check_namedtuple_classdef(defn)
823-
node.node = self.build_namedtuple_typeinfo(
855+
info = self.build_namedtuple_typeinfo(
824856
defn.name, items, types, default_items)
825-
# We only really need the assignments in the body to be type checked later;
826-
# attempting to type check methods may lead to crashes because NamedTuples
827-
# do not have a fully functional TypeInfo.
828-
# TODO remove this hack and add full support for NamedTuple methods
829-
defn.defs.body = [stmt for stmt in defn.defs.body
830-
if isinstance(stmt, AssignmentStmt)]
831-
return True
832-
return False
857+
node.node = info
858+
defn.info = info
859+
return info
860+
return None
833861

834862
def check_namedtuple_classdef(
835863
self, defn: ClassDef) -> Tuple[List[str], List[Type], Dict[str, Expression]]:
836864
NAMEDTUP_CLASS_ERROR = ('Invalid statement in NamedTuple definition; '
837-
'expected "field_name: field_type"')
865+
'expected "field_name: field_type [= default]"')
838866
if self.options.python_version < (3, 6):
839867
self.fail('NamedTuple class syntax is only supported in Python 3.6', defn)
840868
return [], [], {}
@@ -846,10 +874,18 @@ def check_namedtuple_classdef(
846874
for stmt in defn.defs.body:
847875
if not isinstance(stmt, AssignmentStmt):
848876
# Still allow pass or ... (for empty namedtuples).
849-
if (not isinstance(stmt, PassStmt) and
850-
not (isinstance(stmt, ExpressionStmt) and
851-
isinstance(stmt.expr, EllipsisExpr))):
852-
self.fail(NAMEDTUP_CLASS_ERROR, stmt)
877+
if (isinstance(stmt, PassStmt) or
878+
(isinstance(stmt, ExpressionStmt) and
879+
isinstance(stmt.expr, EllipsisExpr))):
880+
continue
881+
# Also allow methods, including decorated ones.
882+
if isinstance(stmt, (Decorator, FuncBase)):
883+
continue
884+
# And docstrings.
885+
if (isinstance(stmt, ExpressionStmt) and
886+
isinstance(stmt.expr, StrExpr)):
887+
continue
888+
self.fail(NAMEDTUP_CLASS_ERROR, stmt)
853889
elif len(stmt.lvalues) > 1 or not isinstance(stmt.lvalues[0], NameExpr):
854890
# An assignment, but an invalid one.
855891
self.fail(NAMEDTUP_CLASS_ERROR, stmt)
@@ -2090,6 +2126,8 @@ def add_field(var: Var, is_initialized_in_class: bool = False,
20902126
add_field(Var('_field_types', dictype), is_initialized_in_class=True)
20912127
add_field(Var('_field_defaults', dictype), is_initialized_in_class=True)
20922128
add_field(Var('_source', strtype), is_initialized_in_class=True)
2129+
add_field(Var('__annotations__', ordereddictype), is_initialized_in_class=True)
2130+
add_field(Var('__doc__', strtype), is_initialized_in_class=True)
20932131

20942132
tvd = TypeVarDef('NT', 1, [], info.tuple_type)
20952133
selftype = TypeVarType(tvd)
@@ -3462,7 +3500,7 @@ def visit_class_def(self, cdef: ClassDef) -> None:
34623500
self.process_nested_classes(cdef)
34633501

34643502
def process_nested_classes(self, outer_def: ClassDef) -> None:
3465-
self.sem.enter_class(outer_def)
3503+
self.sem.enter_class(outer_def.info)
34663504
for node in outer_def.defs.body:
34673505
if isinstance(node, ClassDef):
34683506
node.info = TypeInfo(SymbolTable(), node, self.sem.cur_mod_id)
@@ -3604,8 +3642,11 @@ def visit_func_def(self, fdef: FuncDef) -> None:
36043642
self.errors.pop_function()
36053643

36063644
def visit_class_def(self, tdef: ClassDef) -> None:
3607-
for type in tdef.info.bases:
3608-
self.analyze(type)
3645+
# NamedTuple base classes are validated in check_namedtuple_classdef; we don't have to
3646+
# check them again here.
3647+
if not tdef.info.is_named_tuple:
3648+
for type in tdef.info.bases:
3649+
self.analyze(type)
36093650
# Recompute MRO now that we have analyzed all modules, to pick
36103651
# up superclasses of bases imported from other modules in an
36113652
# import loop. (Only do so if we succeeded the first time.)

0 commit comments

Comments
 (0)