Radically rewrote deep_ast_copy.

jedi/evaluate/helpers.py

@@ -4,18 +4,16 @@ from itertools import chain
 from jedi.parser import tree as pr
 
 
-def deep_ast_copy(obj, new_elements=None):
+def deep_ast_copy(obj, parent=None, new_elements=None):
     """
     Much, much faster than copy.deepcopy, but just for Parser elements (Doesn't
     copy parents).
     """
-    def sort_stmt(key_value):
-        return key_value[0] not in ('_expression_list', '_assignment_details')
 
     if new_elements is None:
         new_elements = {}
 
-    def recursion(obj):
+    def copy_node(obj):
         # If it's already in the cache, just return it.
         try:
             return new_elements[obj]
@@ -24,6 +22,35 @@ def deep_ast_copy(obj, new_elements=None):
             new_obj = copy.copy(obj)
             new_elements[obj] = new_obj
 
+        # Copy children
+        new_children = []
+        for child in obj.children:
+            typ = child.type
+            if typ in ('whitespace', 'operator', 'keyword', 'number', 'string'):
+                # At the moment we're not actually copying those primitive
+                # elements, because there's really no need to. The parents are
+                # obviously wrong, but that's not an issue.
+                new_child = child
+            elif typ == 'name':
+                new_elements[child] = new_child = copy.copy(child)
+                new_child.parent = new_obj
+            else:  # Is a BaseNode.
+                new_child = copy_node(child)
+                new_child.parent = new_obj
+            new_children.append(new_child)
+        new_obj.children = new_children
+
+        # Copy the names_dict (if there is one).
+        try:
+            names_dict = obj.names_dict
+        except AttributeError:
+            pass
+        else:
+            new_obj.names_dict = new_names_dict = {}
+            for string, names in names_dict.items():
+                new_names_dict[string] = [new_elements[n] for n in names]
+        return new_obj
+
         # Gather items
         try:
             items = list(obj.__dict__.items())
@@ -75,7 +102,13 @@ def deep_ast_copy(obj, new_elements=None):
         if isinstance(array_obj, tuple):
             return tuple(copied_array)
         return copied_array
-    return recursion(obj)
+
+    if parent is not None:
+        new_obj = copy_node(obj)
+        for child in new_obj.children:
+            if isinstance(child, (pr.Name, pr.BaseNode)):
+                child.parent = parent
+    return new_obj
 
 
 def call_of_name(name, cut_own_trailer=False):

jedi/evaluate/representation.py

@@ -570,10 +570,7 @@ class FunctionExecution(Executed):
     def __init__(self, evaluator, base, *args, **kwargs):
         super(FunctionExecution, self).__init__(evaluator, base, *args, **kwargs)
         self._copy_dict = {}
-        new_func = helpers.deep_ast_copy(base.base_func, self._copy_dict)
+        new_func = helpers.deep_ast_copy(base.base_func, self, self._copy_dict)
-        for child in new_func.children:
-            if isinstance(child, (pr.Name, pr.BaseNode)):
-                child.parent = self
         self.children = new_func.children
         self.names_dict = new_func.names_dict
 
@@ -653,6 +650,7 @@ class FunctionExecution(Executed):
         return getattr(self.base, name)
 
     def _scope_copy(self, scope):
+        raise NotImplementedError
         """ Copies a scope (e.g. `if foo:`) in an execution """
         if scope != self.base.base_func:
             # Just make sure the parents been copied.
@@ -662,22 +660,22 @@ class FunctionExecution(Executed):
     @common.safe_property
     @memoize_default([])
     def returns(self):
-        return self._copy_list(self.base.returns)
+        return pr.Scope._search_in_scope(self, pr.ReturnStmt)
 
     @common.safe_property
     @memoize_default([])
     def yields(self):
-        return self._copy_list(self.base.yields)
+        return pr.Scope._search_in_scope(self, pr.YieldExpr)
 
     @common.safe_property
     @memoize_default([])
     def statements(self):
-        return self._copy_list(self.base.statements)
+        return pr.Scope._search_in_scope(self, pr.ExprStmt)
 
     @common.safe_property
     @memoize_default([])
     def subscopes(self):
-        return self._copy_list(self.base.subscopes)
+        return pr.Scope._search_in_scope(self, pr.Scope)
 
     def __repr__(self):
         return "<%s of %s>" % (type(self).__name__, self.base)