save assignement nodes and call set_local after walking the tree

--HG--
branch : rebuild

3 files changed, 20 insertions, 13 deletions

diff --git a/_nodes_ast.py b/_nodes_ast.py
index b3e2709f..21340f37 100644
--- a/_nodes_ast.py
+++ b/_nodes_ast.py
@@ -185,11 +185,6 @@ class TreeRebuilder(RebuildVisitor):
         self._delayed['assattr'].append(newnode)
         return newnode
 
-    def visit_assname(self, node):
-        """visit a AssName node by returning a fresh instance of it"""
-        newnode = new.AssName()
-        return newnode
-
     def visit_assert(self, node):
         """visit a Assert node by returning a fresh instance of it"""
         newnode = new.Assert()
@@ -448,9 +443,11 @@ class TreeRebuilder(RebuildVisitor):
         """visit a Name node by returning a fresh instance of it"""
         if self.asscontext == "Del":
             newnode = new.DelName()
+            self._save_assigment(newnode)
         elif self.asscontext is not None: # Ass
             assert self.asscontext == "Ass"
             newnode = new.AssName()
+            self._save_assigment(newnode)
         else:
             newnode = new.Name()
         newnode.name = node.id
diff --git a/_nodes_compiler.py b/_nodes_compiler.py
index f6139fae..9aa50289 100644
--- a/_nodes_compiler.py
+++ b/_nodes_compiler.py
@@ -240,10 +240,12 @@ class TreeRebuilder(RebuildVisitor):
             return delnode
         elif self.asscontext == "Del":
             newnode = self.visit_delname(node)
+            self._save_assigment(newnode)
         elif self.asscontext == "Ass":
             assert node.flags == 'OP_ASSIGN'
             newnode = new.AssName()
             newnode.name = node.name
+            self._save_assigment(newnode)
         return newnode
 
     def visit_assert(self, node):
diff --git a/rebuilder.py b/rebuilder.py
index 49260b50..21118fc3 100644
--- a/rebuilder.py
+++ b/rebuilder.py
@@ -41,7 +41,7 @@ class RebuildVisitor(ASTVisitor):
         self._delayed = dict((name, []) for name in ('class', 'function', 'assattr'))
         self.set_line_info = (ast_mode == '_ast')
         self._ast_mode = (ast_mode == '_ast')
-        self._asscontext = None
+        self._assignments = []
 
     def visit(self, node, parent):
         if node is None: # some attributes of some nodes are just None
@@ -76,6 +76,8 @@ class RebuildVisitor(ASTVisitor):
             delay_method = getattr(self, 'delayed_' + name)
             for node in nodes:
                 delay_method(node)
+        for node, root in self._assignments:
+            self.set_local_name(node, root)
         return newnode
 
     # general visit_<node> methods ############################################
@@ -231,13 +233,19 @@ class RebuildVisitor(ASTVisitor):
         newnode = nodes.Pass()
         return newnode
 
-    def visit_assname(self, node): # XXX parent
-        if self.asscontext is not None:
-            if self._global_names and node.name in self._global_names[-1]:
-                node.root().set_local(node.name, node)
-            else:
-                node.parent.set_local(node.name, node)
-    visit_delname = visit_assname
+    def _save_assigment(self, node):
+        """save assignement situation since node.parent is not available yet"""
+        if self._global_names and node.name in self._global_names[-1]:
+            self._assignments.append((node, True))
+        else:
+            self._assignments.append((node, False))
+
+    def set_local_name(self, node, root):
+        """set local name into the right place"""
+        if root:
+            node.root().set_local(node.name, node)
+        else:
+            node.parent.set_local(node.name, node)
 
     def delayed_assattr(self, node):
         """visit a AssAttr node -> add name to locals, handle members