[pypy-commit] pypy default: Make iterating over generators work (super easy!)

[alex_gaynor]

Author: Alex Gaynor <alex.gay...@gmail.com>
Branch: 
Changeset: r56577:101051858877
Date: 2012-08-04 11:49 -0700
http://bitbucket.org/pypy/pypy/changeset/101051858877/

Log:    Make iterating over generators work (super easy!)

diff --git a/pypy/rpython/test/test_generator.py 
b/pypy/rpython/test/test_generator.py
--- a/pypy/rpython/test/test_generator.py
+++ b/pypy/rpython/test/test_generator.py
@@ -1,3 +1,5 @@
+import py
+
 from pypy.rpython.test.tool import BaseRtypingTest, LLRtypeMixin, OORtypeMixin
 
 
@@ -74,9 +76,24 @@
         res = self.interpret(f, [])
         assert res == 358
 
+    def test_iterating_generator(self):
+        def f():
+            yield 1
+            yield 2
+            yield 3
+        def g():
+            s = 0
+            for x in f():
+                s += x
+            return s
+        res = self.interpret(g, [])
+        assert res == 6
+
 
 class TestLLtype(BaseTestGenerator, LLRtypeMixin):
     pass
 
+
 class TestOOtype(BaseTestGenerator, OORtypeMixin):
-    pass
+    def test_iterating_generator(self):
+        py.test.skip("Iterators aren't supported on OOtype yet")
diff --git a/pypy/translator/generator.py b/pypy/translator/generator.py
--- a/pypy/translator/generator.py
+++ b/pypy/translator/generator.py
@@ -33,8 +33,13 @@
         class Entry(AbstractPosition):
             _immutable_ = True
             varnames = get_variable_names(graph.startblock.inputargs)
+
         def __init__(self, entry):
             self.current = entry
+
+        def __iter__(self):
+            return self
+
     return GeneratorIterator
 
 def replace_graph_with_bootstrap(GeneratorIterator, graph):
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
http://mail.python.org/mailman/listinfo/pypy-commit