[pypy-commit] pypy stmgc-c7: Add TransactionQueue.add_generator()

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: stmgc-c7
Changeset: r76854:2e010e0f1fdc
Date: 2015-04-21 08:36 +0200
http://bitbucket.org/pypy/pypy/changeset/2e010e0f1fdc/

Log:	Add TransactionQueue.add_generator()

diff --git a/lib_pypy/pypy_test/test_transaction.py b/lib_pypy/pypy_test/test_transaction.py
--- a/lib_pypy/pypy_test/test_transaction.py
+++ b/lib_pypy/pypy_test/test_transaction.py
@@ -47,6 +47,19 @@
             print lsts
         assert lsts == (range(10),) * 5, lsts
 
+def test_add_generator():
+    lst = []
+    def do_stuff(n):
+        for i in range(n):
+            lst.append(i)
+            yield
+    tq = transaction.TransactionQueue()
+    tq.add_generator(do_stuff(10))
+    tq.run()
+    if VERBOSE:
+        print lst
+    assert lst == range(10), lst
+
 def test_raise():
     class FooError(Exception):
         pass
diff --git a/lib_pypy/transaction.py b/lib_pypy/transaction.py
--- a/lib_pypy/transaction.py
+++ b/lib_pypy/transaction.py
@@ -128,6 +128,19 @@
         #   thread-local list.
         self._pending.append((f, args, kwds))
 
+    def add_generator(self, generator_iterator):
+        """Register N new transactions to be done by a generator-iterator
+        object.  Each 'yield' marks the limit of transactions.
+        """
+        def transact_until_yield():
+            # Ask for the next item in this transaction.  If we get it,
+            # then the 'for' loop below adds this function again and
+            # returns.
+            for ignored_yielded_value in generator_iterator:
+                self.add(transact_until_yield)
+                return
+        self.add(transact_until_yield)
+
     def run(self, nb_segments=0):
         """Run all transactions, and all transactions started by these
         ones, recursively, until the queue is empty.  If one transaction