| 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]) | |