Support yield inside lambda
authorVitja Makarov <vitja.makarov@gmail.com>
Sun, 9 Jan 2011 16:44:37 +0000 (19:44 +0300)
committerVitja Makarov <vitja.makarov@gmail.com>
Sun, 9 Jan 2011 16:44:37 +0000 (19:44 +0300)
Cython/Compiler/ParseTreeTransforms.py
tests/run/generators.pyx

index 01be309d5283ece533b53e96e9e0dc551036377e..4e759632bf6b0ac0be74555a68ed9e6835bb13e4 100644 (file)
@@ -191,14 +191,20 @@ class PostParse(ScopeTrackingTransform):
         lambda_id = self.lambda_counter
         self.lambda_counter += 1
         node.lambda_name = EncodedString(u'lambda%d' % lambda_id)
-
-        body = Nodes.ReturnStatNode(
-            node.result_expr.pos, value = node.result_expr)
+        collector = YieldNodeCollector()
+        collector.visitchildren(node.result_expr)
+        if collector.yields or isinstance(node.result_expr, ExprNodes.YieldExprNode):
+            body = ExprNodes.YieldExprNode(
+                node.result_expr.pos, arg=node.result_expr)
+            body = Nodes.ExprStatNode(node.result_expr.pos, expr=body)
+        else:
+            body = Nodes.ReturnStatNode(
+                node.result_expr.pos, value=node.result_expr)
         node.def_node = Nodes.DefNode(
             node.pos, name=node.name, lambda_name=node.lambda_name,
             args=node.args, star_arg=node.star_arg,
             starstar_arg=node.starstar_arg,
-            body=body)
+            body=body, doc=None)
         self.visitchildren(node)
         return node
 
@@ -1383,7 +1389,8 @@ class MarkClosureVisitor(CythonTransform):
                                                starstar_arg=node.starstar_arg,
                                                doc=node.doc,
                                                decorators=node.decorators,
-                                               gbody=gbody)
+                                               gbody=gbody,
+                                               lambda_name=node.lambda_name)
             return generator
         return node
 
index 53047284e591cc3602969699da3d867938b7d7a8..af8926c87c2899c49f1303870af4d288d793e3c2 100644 (file)
@@ -287,3 +287,15 @@ def test_nested_yield():
     StopIteration
     """
     yield (yield (yield 1))
+
+def test_inside_lambda():
+    """
+    >>> obj = test_inside_lambda()()
+    >>> next(obj)
+    1
+    >>> obj.send('a')
+    2
+    >>> obj.send('b')
+    ('a', 'b')
+    """
+    return lambda:((yield 1), (yield 2))