Move the option to "exclude" children from the tree traversal from the...

Move the option to "exclude" children from the tree traversal from the "VisitorTransform" class up to its parent "TreeVisitor" class.

Cython/Compiler/FlowControl.py

@@ -661,7 +661,7 @@ class AssignmentCollector(TreeVisitor):
         self.assignments = []
 
     def visit_Node(self):
-        self._visitchildren(self, None)
+        self._visitchildren(self, None, None)
 
     def visit_SingleAssignmentNode(self, node):
         self.assignments.append((node.lhs, node.rhs))

Cython/Compiler/Visitor.pxd

@@ -10,15 +10,14 @@ cdef class TreeVisitor:
     cdef _visit(self, obj)
     cdef find_handler(self, obj)
     cdef _visitchild(self, child, parent, attrname, idx)
-    cdef dict _visitchildren(self, parent, attrs)
-    cpdef visitchildren(self, parent, attrs=*)
+    cdef dict _visitchildren(self, parent, attrs, exclude)
+    cpdef visitchildren(self, parent, attrs=*, exclude=*)
     cdef _raise_compiler_error(self, child, e)
 
 cdef class VisitorTransform(TreeVisitor):
-    cdef dict _process_children(self, parent, attrs=*)
+    cdef dict _process_children(self, parent, attrs=*, exclude=*)
     cpdef visitchildren(self, parent, attrs=*, exclude=*)
     cdef list _flatten_list(self, list orig_list)
-    cdef list _select_attrs(self, attrs, exclude)
     cpdef visitchild(self, parent, str attr, idx=*)
 
 cdef class CythonTransform(VisitorTransform):

Cython/Compiler/Visitor.py

@@ -167,10 +167,12 @@ class TreeVisitor(object):
         raise RuntimeError("Visitor %r does not accept object: %s" % (self, obj))
 
     def visit(self, obj):
+        # generic def entry point for calls from Python subclasses
         return self._visit(obj)
 
     @cython.final
     def _visit(self, obj):
+        # fast cdef entry point for calls from Cython subclasses
         try:
             try:
                 handler_method = self.dispatch_table[type(obj)]
@@ -189,17 +191,20 @@ class TreeVisitor(object):
 
     @cython.final
     def _visitchild(self, child, parent, attrname, idx):
+        # fast cdef entry point for calls from Cython subclasses
         self.access_path.append((parent, attrname, idx))
         result = self._visit(child)
         self.access_path.pop()
         return result
 
-    def visitchildren(self, parent, attrs=None):
-        return self._visitchildren(parent, attrs)
+    def visitchildren(self, parent, attrs=None, exclude=None):
+        # generic def entry point for calls from Python subclasses
+        return self._visitchildren(parent, attrs, exclude)
 
     @cython.final
     @cython.locals(idx=cython.Py_ssize_t)
-    def _visitchildren(self, parent, attrs):
+    def _visitchildren(self, parent, attrs, exclude):
+        # fast cdef entry point for calls from Cython subclasses
         """
         Visits the children of the given parent. If parent is None, returns
         immediately (returning None).
@@ -213,6 +218,7 @@ class TreeVisitor(object):
         result = {}
         for attr in parent.child_attrs:
             if attrs is not None and attr not in attrs: continue
+            if exclude is not None and attr in exclude: continue
             child = getattr(parent, attr)
             if child is not None:
                 if type(child) is list:
@@ -246,18 +252,12 @@ class VisitorTransform(TreeVisitor):
     """
     def visitchildren(self, parent, attrs=None, exclude=None):
         # generic def entry point for calls from Python subclasses
-        if exclude is not None:
-            attrs = self._select_attrs(parent.child_attrs if attrs is None else attrs, exclude)
-        return self._process_children(parent, attrs)
-
-    @cython.final
-    def _select_attrs(self, attrs, exclude):
-        return [name for name in attrs if name not in exclude]
+        return self._process_children(parent, attrs, exclude)
 
     @cython.final
-    def _process_children(self, parent, attrs=None):
+    def _process_children(self, parent, attrs=None, exclude=None):
         # fast cdef entry point for calls from Cython subclasses
-        result = self._visitchildren(parent, attrs)
+        result = self._visitchildren(parent, attrs, exclude)
         for attr, newnode in result.items():
             if type(newnode) is list:
                 newnode = self._flatten_list(newnode)
@@ -756,7 +756,7 @@ class NodeFinder(TreeVisitor):
         elif node is self.node:
             self.found = True
         else:
-            self._visitchildren(node, None)
+            self._visitchildren(node, None, None)
 
 def tree_contains(tree, node):
     finder = NodeFinder(node)