first sketch for handling remaining visits, delays, set_locals etc.

--HG--
branch : rebuild

3 files changed, 48 insertions, 35 deletions

diff --git a/_nodes_ast.py b/_nodes_ast.py
index 9d87718d..c08243d0 100644
--- a/_nodes_ast.py
+++ b/_nodes_ast.py
@@ -174,6 +174,7 @@ class TreeRebuilder(RebuildVisitor):
         newnode.defaults = [self.visit(child, node) for child in node.defaults]
         newnode.vararg = node.vararg
         newnode.kwarg = node.kwarg
+        self._save_argument_name(newnode)
         return newnode
 
     def visit_assattr(self, node):
@@ -277,8 +278,10 @@ class TreeRebuilder(RebuildVisitor):
         """visit a Decorators node by returning a fresh instance of it"""
         newnode = new.Decorators()
         newnode.nodes = [self.visit(child, node) for child in node.decorators]
+        self._delayed['decorators'].append(newnode)
         return newnode
 
+
     def visit_delete(self, node):
         """visit a Delete node by returning a fresh instance of it"""
         newnode = new.Delete()
@@ -335,6 +338,7 @@ class TreeRebuilder(RebuildVisitor):
         """visit a From node by returning a fresh instance of it"""
         names = [(alias.name, alias.asname) for alias in node.names]
         newnode = new.From(node.module, names)
+        self._delayed['from'].append(newnode)
         return newnode
 
     def _visit_function(self, node):
@@ -396,6 +400,7 @@ class TreeRebuilder(RebuildVisitor):
         """visit a Import node by returning a fresh instance of it"""
         newnode = new.Import()
         newnode.names = [(alias.name, alias.asname) for alias in node.names]
+        self._save_import_locals(newnode)
         return newnode
 
     def visit_index(self, node):
@@ -437,6 +442,7 @@ class TreeRebuilder(RebuildVisitor):
         newnode = new.Module()
         _init_set_doc(node, newnode)
         newnode.body = [self.visit(child, node) for child in node.body]
+        newnode.name = node.name
         return newnode
 
     def _visit_name(self, node):
diff --git a/_nodes_compiler.py b/_nodes_compiler.py
index 581ec845..c28bc563 100644
--- a/_nodes_compiler.py
+++ b/_nodes_compiler.py
@@ -163,6 +163,7 @@ def _nodify_args(parent, values):
 
 def args_compiler_to_ast(newnode, node):
     # insert Arguments node
+    # XXX move this to TreeRebuilder.visit_arguments
     if node.flags & 8:
         kwarg = node.argnames.pop()
     else:
@@ -214,8 +215,11 @@ class TreeRebuilder(RebuildVisitor):
 
     def visit_arguments(self, node):
         """visit an Arguments node by returning a fresh instance of it"""
+        XXX
+        # XXX never called ? move args_compiler_to_ast
         newnode = new.Arguments()
         newnode.args = [self.visit(child, node) for child in node.args]
+        self._save_argument_name(newnode)
         newnode.defaults = [self.visit(child, node) for child in node.defaults]
         return newnode
 
@@ -441,6 +445,7 @@ class TreeRebuilder(RebuildVisitor):
     def visit_from(self, node):
         """visit a From node by returning a fresh instance of it"""
         newnode = new.From(node.modname, node.names)
+        self._delayed['from'].append(newnode)
         return newnode
 
 
@@ -506,6 +511,7 @@ class TreeRebuilder(RebuildVisitor):
         """visit an Import node by returning a fresh instance of it"""
         newnode = new.Import()
         newnode.names = node.names
+        self._save_import_locals(newnode)
         return newnode
 
     def visit_index(self, node):
@@ -549,6 +555,7 @@ class TreeRebuilder(RebuildVisitor):
         newnode = new.Module()
         newnode.doc = node.doc
         newnode.body = [self.visit(child, node) for child in node.node.nodes]
+        newnode.name = node.name
         return newnode
 
     def _visit_name(self, node):
diff --git a/rebuilder.py b/rebuilder.py
index 1dc85075..b44342a3 100644
--- a/rebuilder.py
+++ b/rebuilder.py
@@ -50,7 +50,8 @@ class RebuildVisitor(ASTVisitor):
         self.asscontext = None
         self._metaclass = None
         self._global_names = None
-        self._delayed = dict((name, []) for name in ('class', 'function', 'assattr'))
+        self._delayed = dict((name, []) for name in ('class', 'function',
+                              'assattr', 'assign', 'decorators', 'from'))
         self.set_line_info = (ast_mode == '_ast')
         self._ast_mode = (ast_mode == '_ast')
         self._assignments = []
@@ -58,7 +59,8 @@ class RebuildVisitor(ASTVisitor):
     def visit(self, node, parent):
         if node is None: # some attributes of some nodes are just None
             return None
-        # TODO : remove parent: it is never used
+        # TODO : parent is never used; replace by newparent ? Could
+        #        replace set_infos or simplify some other cases ?
         cls_name = node.__class__.__name__
         _method_suffix = REDIRECT.get(cls_name, cls_name).lower()
         _visit = getattr(self, "visit_%s" % _method_suffix )
@@ -96,24 +98,29 @@ class RebuildVisitor(ASTVisitor):
             delay_method = getattr(self, 'delayed_' + name)
             for node in nodes:
                 delay_method(node)
-        for assnode, root in self._assignments:
-            self.set_local_name(assnode, root)
+        for assnode, name, root_local in self._assignments:
+            if root_local:
+                assnode.root().set_local(name, assnode)
+            else:
+                assnode.parent.set_local(name, assnode)
         return newnode
 
+
     # general visit_<node> methods ############################################
 
-    def visit_arguments(self, node): # XXX  parent...
+    def _save_argument_name(self, node):
+        """save argument names for setting locals"""
         if node.vararg:
-            node.parent.set_local(node.vararg, node)
+            self._assignments.append((node, node.vararg, False))
         if node.kwarg:
-            node.parent.set_local(node.kwarg, node)
+            self._assignments.append((node, node.kwarg, False))
 
     def visit_assign(self, node):
         newnode = self._visit_assign(node)
-        # XXX call leave_assign  here ?
+        self._delayed['assign'].append(newnode)
         return newnode
 
-        #def xxx_leave_assign(self, newnode): #XXX  parent...
+    def delayed_assign(self, newnode):
         klass = newnode.parent.frame()
         if (isinstance(klass, nodes.Class)
             and isinstance(newnode.value, nodes.CallFunc)
@@ -125,7 +132,7 @@ class RebuildVisitor(ASTVisitor):
                     if isinstance(meth, nodes.Function):
                         if func_name in ('classmethod', 'staticmethod'):
                             meth.type = func_name
-                        try:
+                        try: # XXX use setdefault ?
                             meth.extra_decorators.append(newnode.value)
                         except AttributeError:
                             meth.extra_decorators = [newnode.value]
@@ -136,8 +143,8 @@ class RebuildVisitor(ASTVisitor):
             self._metaclass[-1] = 'type' # XXX get the actual metaclass
         return newnode
 
-    def visit_class(self, node): # TODO
-        """visit an Class node to become astng"""
+    def visit_class(self, node):
+        """visit a Class node to become astng"""
         newnode = self._visit_class(node)
         newnode.name = node.name
         self._metaclass.append(self._metaclass[-1])
@@ -161,15 +168,16 @@ class RebuildVisitor(ASTVisitor):
         newnode = nodes.Continue()
         return newnode
 
-    def visit_decorators(self, node): # TODO
+    def visit_decorators(self, node):
         """visiting an Decorators node"""
-        return self._visit_decorators(node)
+        newnode = self._visit_decorators(node)
+        self._delayed['decorators'].append(newnode)
+        return newnode
 
-    def leave_decorators(self, node): # XXX parent
+    def delayed_decorators(self, node):
         """python >= 2.4
         visit a Decorator node -> check for classmethod and staticmethod
         """
-        return # TODO
         for decorator_expr in node.nodes:
             if isinstance(decorator_expr, nodes.Name) and \
                    decorator_expr.name in ('classmethod', 'staticmethod'):
@@ -186,7 +194,7 @@ class RebuildVisitor(ASTVisitor):
         newnode = nodes.EmptyNode()
         return newnode
 
-    def visit_from(self, node): # TODO XXX root !
+    def delayed_from(self, node):
         """visit an From node to become astng"""
         # add names imported by the import to locals
         for (name, asname) in node.names:
@@ -206,11 +214,8 @@ class RebuildVisitor(ASTVisitor):
         newnode = self._visit_function(node)
         self._delayed['function'].append(newnode)
         newnode.name = node.name
-        return newnode
-
-    def leave_function(self, node): # TODO
-        """leave a Function node -> pop the last item on the stack"""
         self._global_names.pop()
+        return newnode
 
     def delayed_function(self, newnode):
         frame = newnode.parent.frame()
@@ -229,11 +234,11 @@ class RebuildVisitor(ASTVisitor):
                 self._global_names[-1].setdefault(name, []).append(newnode)
         return newnode
 
-    def visit_import(self, node): # XXX parent !
-        """visit an Import node to become astng"""
-        for (name, asname) in node.names:
+    def _save_import_locals(self, newnode):
+        """save import names to set them in locals later on"""
+        for (name, asname) in newnode.names:
             name = asname or name
-            node.parent.set_local(name.split('.')[0], node)
+            self._assignments.append( (newnode, name.split('.')[0], False) )
 
     def visit_module(self, node):
         """visit an Module node to become astng"""
@@ -253,24 +258,19 @@ class RebuildVisitor(ASTVisitor):
         newnode = nodes.Pass()
         return newnode
 
-    def _save_assigment(self, node):
+    def _save_assigment(self, node, name=None):
         """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)
+            self._assignments.append((node, node.name, True))
         else:
-            node.parent.set_local(node.name, node)
+            self._assignments.append((node, node.name, False))
 
     def delayed_assattr(self, node):
         """visit a AssAttr node -> add name to locals, handle members
         definition
         """
+        # XXX this probably doesn't work as expected : 
+        # while running delayed nodes, not all locals are set, hence bad Inference ?
         try:
             frame = node.frame()
             for infered in node.expr.infer():