Fine-grained: Support NewType and reset subtype caches (python#4656)

NewType work highlighted an issue with subtype caches with stale
information leaking, and this fixes that issue as well. We reset the
subtype cache in two places:

* When calculating the MRO; we reset caches of all base classes as
  well.

* When merging a new version of a TypeInfo, which may have a
  different MRO; we reset all caches of base classes in the old
  MRO, as they might no longer be supertypes.

Author: JukkaL

mypy/checker.py

@@ -874,12 +874,12 @@ def is_trivial_body(self, block: Block) -> bool:
         body = block.body
 
         # Skip a docstring
-        if (isinstance(body[0], ExpressionStmt) and
+        if (body and isinstance(body[0], ExpressionStmt) and
                 isinstance(body[0].expr, (StrExpr, UnicodeExpr))):
             body = block.body[1:]
 
         if len(body) == 0:
-            # There's only a docstring.
+            # There's only a docstring (or no body at all).
             return True
         elif len(body) > 1:
             return False

mypy/nodes.py

@@ -2080,6 +2080,11 @@ def is_cached_subtype_check(self, left: 'mypy.types.Instance',
             return (left, right) in self._cache
         return (left, right) in self._cache_proper
 
+    def reset_subtype_cache(self) -> None:
+        for item in self.mro:
+            item._cache = set()
+            item._cache_proper = set()
+
     def __getitem__(self, name: str) -> 'SymbolTableNode':
         n = self.get(name)
         if n:
@@ -2116,6 +2121,7 @@ def calculate_mro(self) -> None:
         self.is_enum = self._calculate_is_enum()
         # The property of falling back to Any is inherited.
         self.fallback_to_any = any(baseinfo.fallback_to_any for baseinfo in self.mro)
+        self.reset_subtype_cache()
 
     def calculate_metaclass_type(self) -> 'Optional[mypy.types.Instance]':
         declared = self.declared_metaclass

mypy/semanal.py

@@ -2182,7 +2182,7 @@ def build_newtype_typeinfo(self, name: str, old_type: Type, base_type: Instance)
             arg_types=[Instance(info, []), old_type],
             arg_kinds=[arg.kind for arg in args],
             arg_names=['self', 'item'],
-            ret_type=old_type,
+            ret_type=NoneTyp(),
             fallback=self.named_type('__builtins__.function'),
             name=name)
         init_func = FuncDef('__init__', args, Block([]), typ=signature)

mypy/server/astmerge.py

@@ -166,8 +166,8 @@ def visit_overloaded_func_def(self, node: OverloadedFuncDef) -> None:
 
     def visit_class_def(self, node: ClassDef) -> None:
         # TODO additional things?
+        node.info = self.fixup_and_reset_typeinfo(node.info)
         node.defs.body = self.replace_statements(node.defs.body)
-        node.info = self.fixup(node.info)
         info = node.info
         for tv in node.type_vars:
             self.process_type_var_def(tv)
@@ -214,7 +214,7 @@ def visit_ref_expr(self, node: RefExpr) -> None:
 
     def visit_namedtuple_expr(self, node: NamedTupleExpr) -> None:
         super().visit_namedtuple_expr(node)
-        node.info = self.fixup(node.info)
+        node.info = self.fixup_and_reset_typeinfo(node.info)
         self.process_synthetic_type_info(node.info)
 
     def visit_super_expr(self, node: SuperExpr) -> None:
@@ -229,7 +229,7 @@ def visit_call_expr(self, node: CallExpr) -> None:
 
     def visit_newtype_expr(self, node: NewTypeExpr) -> None:
         if node.info:
-            node.info = self.fixup(node.info)
+            node.info = self.fixup_and_reset_typeinfo(node.info)
             self.process_synthetic_type_info(node.info)
         self.fixup_type(node.old_type)
         super().visit_newtype_expr(node)
@@ -240,11 +240,11 @@ def visit_lambda_expr(self, node: LambdaExpr) -> None:
 
     def visit_typeddict_expr(self, node: TypedDictExpr) -> None:
         super().visit_typeddict_expr(node)
-        node.info = self.fixup(node.info)
+        node.info = self.fixup_and_reset_typeinfo(node.info)
         self.process_synthetic_type_info(node.info)
 
     def visit_enum_call_expr(self, node: EnumCallExpr) -> None:
-        node.info = self.fixup(node.info)
+        node.info = self.fixup_and_reset_typeinfo(node.info)
         self.process_synthetic_type_info(node.info)
         super().visit_enum_call_expr(node)
 
@@ -269,6 +269,19 @@ def fixup(self, node: SN) -> SN:
             return cast(SN, new)
         return node
 
+    def fixup_and_reset_typeinfo(self, node: TypeInfo) -> TypeInfo:
+        """Fix-up type info and reset subtype caches.
+
+        This needs to be called at least once per each merged TypeInfo, as otherwise we
+        may leak stale caches.
+        """
+        if node in self.replacements:
+            # The subclass relationships may change, so reset all caches relevant to the
+            # old MRO.
+            new = cast(TypeInfo, self.replacements[node])
+            new.reset_subtype_cache()
+        return self.fixup(node)
+
     def fixup_type(self, typ: Optional[Type]) -> None:
         if typ is not None:
             typ.accept(TypeReplaceVisitor(self.replacements))

mypy/server/deps.py

@@ -92,7 +92,7 @@ class 'mod.Cls'. This can also refer to an attribute inherited from a
     ComparisonExpr, GeneratorExpr, DictionaryComprehension, StarExpr, PrintStmt, ForStmt, WithStmt,
     TupleExpr, ListExpr, OperatorAssignmentStmt, DelStmt, YieldFromExpr, Decorator, Block,
     TypeInfo, FuncBase, OverloadedFuncDef, RefExpr, SuperExpr, Var, NamedTupleExpr, TypedDictExpr,
-    LDEF, MDEF, GDEF, FuncItem, TypeAliasExpr,
+    LDEF, MDEF, GDEF, FuncItem, TypeAliasExpr, NewTypeExpr,
     op_methods, reverse_op_methods, ops_with_inplace_method, unary_op_methods
 )
 from mypy.traverser import TraverserVisitor
@@ -211,18 +211,27 @@ def visit_class_def(self, o: ClassDef) -> None:
         # Add dependencies to type variables of a generic class.
         for tv in o.type_vars:
             self.add_dependency(make_trigger(tv.fullname), target)
-        # Add dependencies to base types.
-        for base in o.info.bases:
+        self.process_type_info(o.info)
+        super().visit_class_def(o)
+        self.is_class = old_is_class
+        self.scope.leave()
+
+    def visit_newtype_expr(self, o: NewTypeExpr) -> None:
+        if o.info:
+            self.scope.enter_class(o.info)
+            self.process_type_info(o.info)
+            self.scope.leave()
+
+    def process_type_info(self, info: TypeInfo) -> None:
+        target = self.scope.current_full_target()
+        for base in info.bases:
             self.add_type_dependencies(base, target=target)
-        if o.info.tuple_type:
-            self.add_type_dependencies(o.info.tuple_type, target=make_trigger(target))
-        if o.info.typeddict_type:
-            self.add_type_dependencies(o.info.typeddict_type, target=make_trigger(target))
+        if info.tuple_type:
+            self.add_type_dependencies(info.tuple_type, target=make_trigger(target))
+        if info.typeddict_type:
+            self.add_type_dependencies(info.typeddict_type, target=make_trigger(target))
         # TODO: Add dependencies based on remaining TypeInfo attributes.
-        super().visit_class_def(o)
         self.add_type_alias_deps(self.scope.current_target())
-        self.is_class = old_is_class
-        info = o.info
         for name, node in info.names.items():
             if isinstance(node.node, Var):
                 for base_info in non_trivial_bases(info):
@@ -236,7 +245,6 @@ def visit_class_def(self, o: ClassDef) -> None:
                                     target=make_trigger(info.fullname() + '.' + name))
             self.add_dependency(make_trigger(base_info.fullname() + '.__init__'),
                                 target=make_trigger(info.fullname() + '.__init__'))
-        self.scope.leave()
 
     def visit_import(self, o: Import) -> None:
         for id, as_id in o.ids:

test-data/unit/deps-statements.test

@@ -655,3 +655,20 @@ class C:
 <m.C> -> m.C
 <sys.platform> -> m
 <sys> -> m
+
+[case testNewType]
+from typing import NewType
+from m import C
+
+N = NewType('N', C)
+
+def f(n: N) -> None:
+    pass
+[file m.py]
+class C:
+    x: int
+[out]
+<m.N> -> <m.f>, m, m.f
+<m.C.__init__> -> <m.N.__init__>
+<m.C.x> -> <m.N.x>
+<m.C> -> m, m.N

test-data/unit/diff.test

@@ -681,6 +681,30 @@ B = Dict[str, S]
 __main__.A
 __main__.T
 
+[case testNewType]
+from typing import NewType
+class C: pass
+class D: pass
+N1 = NewType('N1', C)
+N2 = NewType('N2', D)
+N3 = NewType('N3', C)
+class N4(C): pass
+[file next.py]
+from typing import NewType
+class C: pass
+class D(C): pass
+N1 = NewType('N1', C)
+N2 = NewType('N2', D)
+class N3(C): pass
+N4 = NewType('N4', C)
+[out]
+__main__.D
+__main__.N2
+__main__.N3
+__main__.N3.__init__
+__main__.N4
+__main__.N4.__init__
+
 [case testChangeGenericBaseClassOnly]
 from typing import List
 class C(List[int]): pass

test-data/unit/fine-grained.test

@@ -1528,7 +1528,8 @@ import a
 [file a.py]
 from typing import Dict, NewType
 
-N = NewType('N', int)
+class A: pass
+N = NewType('N', A)
 
 a: Dict[N, int]
 
@@ -1538,7 +1539,8 @@ def f(self, x: N) -> None:
 [file a.py.2]
 from typing import Dict, NewType  # dummy change
 
-N = NewType('N', int)
+class A: pass
+N = NewType('N', A)
 
 a: Dict[N, int]
 
@@ -2498,6 +2500,71 @@ else:
 [out]
 ==
 
+[case testNewTypeDependencies1]
+from a import N
+
+def f(x: N) -> None:
+    x.y = 1
+[file a.py]
+from typing import NewType
+from b import C
+
+N = NewType('N', C)
+[file b.py]
+class C:
+    y: int
+[file b.py.2]
+class C:
+    y: str
+[out]
+==
+main:4: error: Incompatible types in assignment (expression has type "int", variable has type "str")
+
+[case testNewTypeDependencies2]
+from a import N
+from b import C, D
+
+def f(x: C) -> None: pass
+
+def g(x: N) -> None:
+    f(x)
+[file a.py]
+from typing import NewType
+from b import D
+
+N = NewType('N', D)
+[file b.py]
+class C: pass
+class D(C): pass
+[file b.py.2]
+class C: pass
+class D: pass
+[out]
+==
+main:7: error: Argument 1 to "f" has incompatible type "N"; expected "C"
+
+[case testNewTypeDependencies3]
+from a import N
+
+def f(x: N) -> None:
+    x.y
+[file a.py]
+from typing import NewType
+from b import C
+N = NewType('N', C)
+[file a.py.2]
+from typing import NewType
+from b import D
+N = NewType('N', D)
+[file b.py]
+class C:
+    y: int
+class D:
+    pass
+[out]
+==
+main:4: error: "N" has no attribute "y"
+
 [case testNamedTupleWithinFunction]
 from typing import NamedTuple
 import b