File size: 1,240 Bytes
db4a26f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
from greenlet import greenlet
from . import TestCase
class genlet(greenlet):
parent = None
def __init__(self, *args, **kwds):
self.args = args
self.kwds = kwds
def run(self):
fn, = self.fn
fn(*self.args, **self.kwds)
def __iter__(self):
return self
def __next__(self):
self.parent = greenlet.getcurrent()
result = self.switch()
if self:
return result
raise StopIteration
next = __next__
def Yield(value):
g = greenlet.getcurrent()
while not isinstance(g, genlet):
if g is None:
raise RuntimeError('yield outside a genlet')
g = g.parent
g.parent.switch(value)
def generator(func):
class Generator(genlet):
fn = (func,)
return Generator
# ____________________________________________________________
class GeneratorTests(TestCase):
def test_generator(self):
seen = []
def g(n):
for i in range(n):
seen.append(i)
Yield(i)
g = generator(g)
for _ in range(3):
for j in g(5):
seen.append(j)
self.assertEqual(seen, 3 * [0, 0, 1, 1, 2, 2, 3, 3, 4, 4])
|