Repository: python/mypy

Issue:
New semantic analyzer: crash in AST strip
This daemon test case crashes with a KeyError:

```
[case testStripNewAnalyzer]
# flags: --ignore-missing-imports
import a
[file a.py]
from typing import List
from b import B

class A:
    def __init__(self) -> None:
        self.x: List[int] = []

    def method(self) -> None:
        B()
        self.x = []

[file b.py]
class B: ...
[delete b.py.2]

[builtins fixtures/list.pyi]
[out]
==
```

Here's the traceback:

```
data: /Users/jukka/src/mypy/test-data/unit/fine-grained.test:8932:
/Users/jukka/src/mypy/mypy/test/testfinegrained.py:122: in run_case
    new_messages = self.run_check(server, sources)
/Users/jukka/src/mypy/mypy/test/testfinegrained.py:194: in run_check
    response = server.check(sources)
/Users/jukka/src/mypy/mypy/dmypy_server.py:366: in check
    res = self.fine_grained_increment(sources)
/Users/jukka/src/mypy/mypy/dmypy_server.py:463: in fine_grained_increment
    messages = self.fine_grained_manager.update(changed, removed)
/Users/jukka/src/mypy/mypy/server/update.py:241: in update
    result = self.update_one(changed_modules, initial_set, removed_set)
/Users/jukka/src/mypy/mypy/server/update.py:307: in update_one
    result = self.update_module(next_id, next_path, next_id in removed_set)
/Users/jukka/src/mypy/mypy/server/update.py:389: in update_module
    targets_with_errors=set(), processed_targets=self.processed_targets)
/Users/jukka/src/mypy/mypy/server/update.py:835: in propagate_changes_using_dependencies
    triggered |= reprocess_nodes(manager, graph, id, nodes, deps, processed_targets)
/Users/jukka/src/mypy/mypy/server/update.py:953: in reprocess_nodes
    new_patches = strip_target_new(deferred.node)
/Users/jukka/src/mypy/mypy/server/aststripnew.py:35: in strip_target_new
    node.accept(visitor)
/Users/jukka/src/mypy/mypy/nodes.py:655: in accept
    return visitor.visit_func_def(self)
/Users/jukka/src/mypy/mypy/server/aststripnew.py:122: in visit_func_def
    super().visit_func_def(node)
/Users/jukka/src/mypy/mypy/traverser.py:52: in visit_func_def
    self.visit_func(o)
/Users/jukka/src/mypy/mypy/traverser.py:49: in visit_func
    o.body.accept(self)
/Users/jukka/src/mypy/mypy/nodes.py:978: in accept
    return visitor.visit_block(self)
/Users/jukka/src/mypy/mypy/server/aststripnew.py:60: in visit_block
    super().visit_block(b)
/Users/jukka/src/mypy/mypy/traverser.py:37: in visit_block
    s.accept(self)
/Users/jukka/src/mypy/mypy/nodes.py:1036: in accept
    return visitor.visit_assignment_stmt(self)
/Users/jukka/src/mypy/mypy/server/aststripnew.py:152: in visit_assignment_stmt
    self.process_lvalue_in_method(lvalue)
/Users/jukka/src/mypy/mypy/server/aststripnew.py:203: in process_lvalue_in_method
    del self.type.names[lvalue.name]
E   KeyError: 'x'
```


Discussion:

Resolving patch:
diff --git a/mypy/newsemanal/semanal_main.py b/mypy/newsemanal/semanal_main.py
--- a/mypy/newsemanal/semanal_main.py
+++ b/mypy/newsemanal/semanal_main.py
@@ -28,7 +28,7 @@
 from typing import List, Tuple, Optional, Union, Callable, Iterator
 
 from mypy.nodes import (
-    MypyFile, TypeInfo, FuncDef, Decorator, OverloadedFuncDef
+    MypyFile, TypeInfo, FuncDef, Decorator, OverloadedFuncDef, SymbolTableNode, Var, ClassDef
 )
 from mypy.newsemanal.semanal_typeargs import TypeArgumentAnalyzer
 from mypy.state import strict_optional_set
@@ -42,6 +42,7 @@
 from mypy.errors import Errors
 from mypy.newsemanal.semanal_infer import infer_decorator_signature_if_simple
 from mypy.checker import FineGrainedDeferredNode
+from mypy.server.aststripnew import SavedAttributes
 import mypy.build
 
 MYPY = False
@@ -96,21 +97,26 @@ def cleanup_builtin_scc(state: 'State') -> None:
     remove_imported_names_from_symtable(state.tree.names, 'builtins')
 
 
-def semantic_analysis_for_targets(state: 'State', nodes: List[FineGrainedDeferredNode],
-                                  graph: 'Graph', strip_patches: List[Callable[[], None]]) -> None:
+def semantic_analysis_for_targets(
+        state: 'State',
+        nodes: List[FineGrainedDeferredNode],
+        graph: 'Graph',
+        saved_attrs: SavedAttributes) -> None:
     """Semantically analyze only selected nodes in a given module.
 
     This essentially mirrors the logic of semantic_analysis_for_scc()
     except that we process only some targets. This is used in fine grained
     incremental mode, when propagating an update.
 
-    The strip_patches are additional patches that may be produced by aststrip.py to
-    re-introduce implicitly declared instance variables (attributes defined on self).
+    The saved_attrs are implicitly declared instance attributes (attributes
+    defined on self) removed by AST stripper that may need to be reintroduced
+    here.  They must be added before any methods are analyzed.
     """
     patches = []  # type: Patches
     if any(isinstance(n.node, MypyFile) for n in nodes):
         # Process module top level first (if needed).
         process_top_levels(graph, [state.id], patches)
+    restore_saved_attrs(saved_attrs)
     analyzer = state.manager.new_semantic_analyzer
     for n in nodes:
         if isinstance(n.node, MypyFile):
@@ -119,13 +125,30 @@ def semantic_analysis_for_targets(state: 'State', nodes: List[FineGrainedDeferre
         process_top_level_function(analyzer, state, state.id,
                                    n.node.fullname(), n.node, n.active_typeinfo, patches)
     apply_semantic_analyzer_patches(patches)
-    for patch in strip_patches:
-        patch()
 
     check_type_arguments_in_targets(nodes, state, state.manager.errors)
     calculate_class_properties(graph, [state.id], state.manager.errors)
 
 
+def restore_saved_attrs(saved_attrs: SavedAttributes) -> None:
+    """Restore instance variables removed during AST strip that haven't been added yet."""
+    for (cdef, name), sym in saved_attrs.items():
+        info = cdef.info
+        existing = info.get(name)
+        defined_in_this_class = name in info.names
+        assert isinstance(sym.node, Var)
+        # This needs to mimic the logic in SemanticAnalyzer.analyze_member_lvalue()
+        # regarding the existing variable in class body or in a superclass:
+        # If the attribute of self is not defined in superclasses, create a new Var.
+        if (existing is None or
+                # (An abstract Var is considered as not defined.)
+                (isinstance(existing.node, Var) and existing.node.is_abstract_var) or
+                # Also an explicit declaration on self creates a new Var unless
+                # there is already one defined in the class body.
+                sym.node.explicit_self_type and not defined_in_this_class):
+            info.names[name] = sym
+
+
 def process_top_levels(graph: 'Graph', scc: List[str], patches: Patches) -> None:
     # Process top levels until everything has been bound.
 
diff --git a/mypy/server/aststripnew.py b/mypy/server/aststripnew.py
--- a/mypy/server/aststripnew.py
+++ b/mypy/server/aststripnew.py
@@ -4,7 +4,7 @@
 """
 
 import contextlib
-from typing import Union, Iterator, Optional, List, Callable
+from typing import Union, Iterator, Optional, List, Callable, Dict, Tuple
 
 from mypy.nodes import (
     FuncDef, NameExpr, MemberExpr, RefExpr, MypyFile, ClassDef, AssignmentStmt,
@@ -18,8 +18,11 @@
 from mypy.server.aststrip import nothing
 
 
-def strip_target_new(node: Union[MypyFile, FuncDef, OverloadedFuncDef]
-                     ) -> List[Callable[[], None]]:
+SavedAttributes = Dict[Tuple[ClassDef, str], SymbolTableNode]
+
+
+def strip_target_new(node: Union[MypyFile, FuncDef, OverloadedFuncDef],
+                     saved_attrs: SavedAttributes) -> None:
     """Reset a fine-grained incremental target to state before main pass of semantic analysis.
 
     The most notable difference from the old version of strip_target() is that new semantic
@@ -28,16 +31,15 @@ def strip_target_new(node: Union[MypyFile, FuncDef, OverloadedFuncDef]
     defined as attributes on self. This is done by patches (callbacks) returned from this function
     that re-add these variables when called.
     """
-    visitor = NodeStripVisitor()
+    visitor = NodeStripVisitor(saved_attrs)
     if isinstance(node, MypyFile):
         visitor.strip_file_top_level(node)
     else:
         node.accept(visitor)
-    return visitor.patches
 
 
 class NodeStripVisitor(TraverserVisitor):
-    def __init__(self) -> None:
+    def __init__(self, saved_class_attrs: SavedAttributes) -> None:
         # The current active class.
         self.type = None  # type: Optional[TypeInfo]
         # This is True at class scope, but not in methods.
@@ -45,7 +47,12 @@ def __init__(self) -> None:
         # By default, process function definitions. If False, don't -- this is used for
         # processing module top levels.
         self.recurse_into_functions = True
-        self.patches = []  # type: List[Callable[[], None]]
+        # These attributes were removed from top-level classes during strip and
+        # will be added afterwards (if no existing definition is found). These
+        # must be added back before semantically analyzing any methods. These
+        # allow moving attribute definition from a method (through self.x) to a
+        # definition inside class body (x = ...).
+        self.saved_class_attrs = saved_class_attrs
 
     def strip_file_top_level(self, file_node: MypyFile) -> None:
         """Strip a module top-level (don't recursive into functions)."""
@@ -66,7 +73,7 @@ def visit_class_def(self, node: ClassDef) -> None:
         # be lost if we only reprocess top-levels (this kills TypeInfos)
         # but not the methods that defined those variables.
         if not self.recurse_into_functions:
-            self.prepare_implicit_var_patches(node)
+            self.save_implicit_attributes(node)
         # We need to delete any entries that were generated by plugins,
         # since they will get regenerated.
         to_delete = {v.node for v in node.info.names.values() if v.plugin_generated}
@@ -81,30 +88,11 @@ def visit_class_def(self, node: ClassDef) -> None:
         # Kill the TypeInfo, since there is none before semantic analysis.
         node.info = CLASSDEF_NO_INFO
 
-    def prepare_implicit_var_patches(self, node: ClassDef) -> None:
+    def save_implicit_attributes(self, node: ClassDef) -> None:
         """Produce callbacks that re-add attributes defined on self."""
         for name, sym in node.info.names.items():
             if isinstance(sym.node, Var) and sym.implicit:
-                explicit_self_type = sym.node.explicit_self_type
-
-                def patch(name: str, sym: SymbolTableNode) -> None:
-                    existing = node.info.get(name)
-                    defined_in_this_class = name in node.info.names
-                    # This needs to mimic the logic in SemanticAnalyzer.analyze_member_lvalue()
-                    # regarding the existing variable in class body or in a superclass:
-                    # If the attribute of self is not defined in superclasses, create a new Var.
-                    if (existing is None or
-                            # (An abstract Var is considered as not defined.)
-                            (isinstance(existing.node, Var) and existing.node.is_abstract_var) or
-                            # Also an explicit declaration on self creates a new Var unless
-                            # there is already one defined in the class body.
-                            explicit_self_type and not defined_in_this_class):
-                        node.info.names[name] = sym
-
-                # Capture the current name, sym in a weird hacky way,
-                # because mypyc doesn't yet support capturing them in
-                # the usual hacky way (as default arguments).
-                self.patches.append((lambda name, sym: lambda: patch(name, sym))(name, sym))
+                self.saved_class_attrs[node, name] = sym
 
     def visit_func_def(self, node: FuncDef) -> None:
         if not self.recurse_into_functions:
@@ -199,6 +187,9 @@ def process_lvalue_in_method(self, lvalue: Node) -> None:
                 # self, since only those can define new attributes.
                 assert self.type is not None
                 del self.type.names[lvalue.name]
+                key = (self.type.defn, lvalue.name)
+                if key in self.saved_class_attrs:
+                    del self.saved_class_attrs[key]
         elif isinstance(lvalue, (TupleExpr, ListExpr)):
             for item in lvalue.items:
                 self.process_lvalue_in_method(item)
diff --git a/mypy/server/update.py b/mypy/server/update.py
--- a/mypy/server/update.py
+++ b/mypy/server/update.py
@@ -139,7 +139,7 @@
 from mypy.newsemanal.semanal_main import semantic_analysis_for_scc, semantic_analysis_for_targets
 from mypy.server.astmerge import merge_asts
 from mypy.server.aststrip import strip_target
-from mypy.server.aststripnew import strip_target_new
+from mypy.server.aststripnew import strip_target_new, SavedAttributes
 from mypy.server.deps import get_dependencies_of_target
 from mypy.server.target import module_prefix, split_target, trigger_to_target
 from mypy.server.trigger import make_trigger, WILDCARD_TAG
@@ -944,18 +944,17 @@ def key(node: FineGrainedDeferredNode) -> int:
     manager.errors.clear_errors_in_targets(file_node.path, targets)
 
     # Strip semantic analysis information.
-    patches = []  # type: List[Callable[[], None]]
+    saved_attrs = {}  # type: SavedAttributes
     for deferred in nodes:
         processed_targets.append(deferred.node.fullname())
         if not manager.options.new_semantic_analyzer:
             strip_target(deferred.node)
         else:
-            new_patches = strip_target_new(deferred.node)
-            patches.extend(new_patches)
+            strip_target_new(deferred.node, saved_attrs)
     if not options.new_semantic_analyzer:
         re_analyze_nodes(file_node, nodes, manager, options)
     else:
-        semantic_analysis_for_targets(graph[module_id], nodes, graph, patches)
+        semantic_analysis_for_targets(graph[module_id], nodes, graph, saved_attrs)
     # Merge symbol tables to preserve identities of AST nodes. The file node will remain
     # the same, but other nodes may have been recreated with different identities, such as
     # NamedTuples defined using assignment statements.


Test patch:
diff --git a/test-data/unit/fine-grained.test b/test-data/unit/fine-grained.test
--- a/test-data/unit/fine-grained.test
+++ b/test-data/unit/fine-grained.test
@@ -8928,3 +8928,26 @@ class A(Generic[T]): ...
 [out]
 ==
 main:5: error: Missing type parameters for generic type
+
+[case testStripNewAnalyzer]
+# flags: --ignore-missing-imports
+import a
+[file a.py]
+from typing import List
+from b import B
+
+class A:
+    def __init__(self) -> None:
+        self.x: List[int] = []
+
+    def method(self) -> None:
+        B()
+        self.x = []
+
+[file b.py]
+class B: ...
+[delete b.py.2]
+
+[builtins fixtures/list.pyi]
+[out]
+==