JezK
Edit File: test_generator_nested.py
from greenlet import greenlet from . import TestCase from .leakcheck import fails_leakcheck class genlet(greenlet): parent = None def __init__(self, *args, **kwds): self.args = args self.kwds = kwds self.child = None def run(self): # Note the function is packed in a tuple # to avoid creating a bound method for it. fn, = self.fn fn(*self.args, **self.kwds) def __iter__(self): return self def set_child(self, child): self.child = child def __next__(self): if self.child: child = self.child while child.child: tmp = child child = child.child tmp.child = None result = child.switch() else: self.parent = greenlet.getcurrent() result = self.switch() if self: return result raise StopIteration next = __next__ def Yield(value, level=1): g = greenlet.getcurrent() while level != 0: if not isinstance(g, genlet): raise RuntimeError('yield outside a genlet') if level > 1: g.parent.set_child(g) g = g.parent level -= 1 g.switch(value) def Genlet(func): class TheGenlet(genlet): fn = (func,) return TheGenlet # ____________________________________________________________ def g1(n, seen): for i in range(n): seen.append(i + 1) yield i def g2(n, seen): for i in range(n): seen.append(i + 1) Yield(i) g2 = Genlet(g2) def nested(i): Yield(i) def g3(n, seen): for i in range(n): seen.append(i + 1) nested(i) g3 = Genlet(g3) def a(n): if n == 0: return for ii in ax(n - 1): Yield(ii) Yield(n) ax = Genlet(a) def perms(l): if len(l) > 1: for e in l: # No syntactical sugar for generator expressions x = [Yield([e] + p) for p in perms([x for x in l if x != e])] assert x else: Yield(l) perms = Genlet(perms) def gr1(n): for ii in range(1, n): Yield(ii) Yield(ii * ii, 2) gr1 = Genlet(gr1) def gr2(n, seen): for ii in gr1(n): seen.append(ii) gr2 = Genlet(gr2) class NestedGeneratorTests(TestCase): def test_layered_genlets(self): seen = [] for ii in gr2(5, seen): seen.append(ii) self.assertEqual(seen, [1, 1, 2, 4, 3, 9, 4, 16]) @fails_leakcheck def test_permutations(self): gen_perms = perms(list(range(4))) permutations = list(gen_perms) self.assertEqual(len(permutations), 4 * 3 * 2 * 1) self.assertIn([0, 1, 2, 3], permutations) self.assertIn([3, 2, 1, 0], permutations) res = [] for ii in zip(perms(list(range(4))), perms(list(range(3)))): res.append(ii) self.assertEqual( res, [([0, 1, 2, 3], [0, 1, 2]), ([0, 1, 3, 2], [0, 2, 1]), ([0, 2, 1, 3], [1, 0, 2]), ([0, 2, 3, 1], [1, 2, 0]), ([0, 3, 1, 2], [2, 0, 1]), ([0, 3, 2, 1], [2, 1, 0])]) # XXX Test to make sure we are working as a generator expression def test_genlet_simple(self): for g in [g1, g2, g3]: seen = [] for _ in range(3): for j in g(5, seen): seen.append(j) self.assertEqual(seen, 3 * [1, 0, 2, 1, 3, 2, 4, 3, 5, 4]) def test_genlet_bad(self): try: Yield(10) except RuntimeError: pass def test_nested_genlets(self): seen = [] for ii in ax(5): seen.append(ii)