WIP: Rewrite references to inner functions in treetransform

Author: rwbarton

mypy/treetransform.py

@@ -48,6 +48,7 @@ def __init__(self) -> None:
         # There may be multiple references to a Var node. Keep track of
         # Var translations using a dictionary.
         self.var_map = {}  # type: Dict[Var, Var]
+        self.func_map = {}  # type: Dict[FuncDef, FuncDef]
 
     def visit_mypy_file(self, node: MypyFile) -> Node:
         # NOTE: The 'names' and 'imports' instance variables will be empty!
@@ -98,6 +99,12 @@ def copy_argument(self, argument: Argument) -> Argument:
 
     def visit_func_def(self, node: FuncDef) -> FuncDef:
         # Note that a FuncDef must be transformed to a FuncDef.
+
+        # These contortions are needed to handle the case of recursive
+        # references inside the function being transformed.
+        result = FuncDef(node.name(), node.arguments, node.body, None)
+        self.func_map[node] = result
+
         new = FuncDef(node.name(),
                       [self.copy_argument(arg) for arg in node.arguments],
                       self.block(node.body),
@@ -113,7 +120,8 @@ def visit_func_def(self, node: FuncDef) -> FuncDef:
         new.is_class = node.is_class
         new.is_property = node.is_property
         new.original_def = node.original_def
-        return new
+        result.__dict__ = new.__dict__
+        return result
 
     def visit_func_expr(self, node: FuncExpr) -> Node:
         new = FuncExpr([self.copy_argument(arg) for arg in node.arguments],
@@ -332,6 +340,9 @@ def copy_ref(self, new: RefExpr, original: RefExpr) -> None:
         target = original.node
         if isinstance(target, Var):
             target = self.visit_var(target)
+        elif isinstance(target, FuncDef):
+            if target in self.func_map:
+                target = self.func_map[target]
         new.node = target
         new.is_def = original.is_def