Merge MarkClosureVisitor and MarkGeneratorVisitor

Cython/Compiler/Main.py

@@ -101,7 +101,6 @@ class Context(object):
         from ParseTreeTransforms import WithTransform, NormalizeTree, PostParse, PxdPostParse
         from ParseTreeTransforms import AnalyseDeclarationsTransform, AnalyseExpressionsTransform
         from ParseTreeTransforms import CreateClosureClasses, MarkClosureVisitor, DecoratorTransform
-        from ParseTreeTransforms import MarkGeneratorVisitor
         from ParseTreeTransforms import InterpretCompilerDirectives, TransformBuiltinMethods
         from ParseTreeTransforms import ExpandInplaceOperators
         from TypeInference import MarkAssignments, MarkOverflowingArithmetic
@@ -134,7 +133,6 @@ class Context(object):
             InterpretCompilerDirectives(self, self.compiler_directives),
             _align_function_definitions,
             MarkClosureVisitor(self),
-            MarkGeneratorVisitor(self),
             ConstantFolding(),
             FlattenInListTransform(),
             WithTransform(self),

Cython/Compiler/ParseTreeTransforms.py

@@ -1301,39 +1301,7 @@ class AlignFunctionDefinitions(CythonTransform):
         # Enable this when internal def functions are allowed. 
         # self.visitchildren(node)
         return node
-        
-
-class MarkClosureVisitor(CythonTransform):
-
-    def visit_ModuleNode(self, node):
-        self.needs_closure = False
-        self.visitchildren(node)
-        return node
 
-    def visit_FuncDefNode(self, node):
-        self.needs_closure = False
-        self.visitchildren(node)
-        node.needs_closure = self.needs_closure
-        self.needs_closure = True
-        return node
-
-    def visit_CFuncDefNode(self, node):
-        self.visit_FuncDefNode(node)
-        if node.needs_closure:
-            error(node.pos, "closures inside cdef functions not yet supported")
-        return node
-
-    def visit_LambdaNode(self, node):
-        self.needs_closure = False
-        self.visitchildren(node)
-        node.needs_closure = self.needs_closure
-        self.needs_closure = True
-        return node
-
-    def visit_ClassDefNode(self, node):
-        self.visitchildren(node)
-        self.needs_closure = True
-        return node
 
 class ClosureTempAllocator(object):
     def __init__(self, klass=None):
@@ -1400,23 +1368,21 @@ class YieldNodeCollector(TreeVisitor):
     def visit_DefNode(self, node):
         pass
 
-class MarkGeneratorVisitor(CythonTransform):
-    """XXX: merge me with MarkClosureVisitor"""
-    def __init__(self, context):
-        super(MarkGeneratorVisitor, self).__init__(context)
+class MarkClosureVisitor(CythonTransform):
 
     def visit_ModuleNode(self, node):
+        self.needs_closure = False
         self.visitchildren(node)
         return node
 
-    def visit_ClassDefNode(self, node):
+    def visit_FuncDefNode(self, node):
+        self.needs_closure = False
         self.visitchildren(node)
-        return node
+        node.needs_closure = self.needs_closure
+        self.needs_closure = True
 
-    def visit_FuncDefNode(self, node):
         collector = YieldNodeCollector()
         collector.visitchildren(node)
-        self.visitchildren(node)
 
         if collector.yields:
             allocator = ClosureTempAllocator()
@@ -1429,6 +1395,24 @@ class MarkGeneratorVisitor(CythonTransform):
             node.yields = collector.yields
         return node
 
+    def visit_CFuncDefNode(self, node):
+        self.visit_FuncDefNode(node)
+        if node.needs_closure:
+            error(node.pos, "closures inside cdef functions not yet supported")
+        return node
+
+    def visit_LambdaNode(self, node):
+        self.needs_closure = False
+        self.visitchildren(node)
+        node.needs_closure = self.needs_closure
+        self.needs_closure = True
+        return node
+
+    def visit_ClassDefNode(self, node):
+        self.visitchildren(node)
+        self.needs_closure = True
+        return node
+
 
 class CreateClosureClasses(CythonTransform):
     # Output closure classes in module scope for all functions