codekingpro/portable-devtools

1
2from greenlet import greenlet
3from . import TestCase
4from .leakcheck import fails_leakcheck
5
6class genlet(greenlet):
7    parent = None
8    def __init__(self, *args, **kwds):
9        self.args = args
10        self.kwds = kwds
11        self.child = None
12
13    def run(self):
14        # Note the function is packed in a tuple
15        # to avoid creating a bound method for it.
16        fn, = self.fn
17        fn(*self.args, **self.kwds)
18
19    def __iter__(self):
20        return self
21
22    def set_child(self, child):
23        self.child = child
24
25    def __next__(self):
26        if self.child:
27            child = self.child
28            while child.child:
29                tmp = child
30                child = child.child
31                tmp.child = None
32
33            result = child.switch()
34        else:
35            self.parent = greenlet.getcurrent()
36            result = self.switch()
37
38        if self:
39            return result
40
41        raise StopIteration
42
43    next = __next__
44
45def Yield(value, level=1):
46    g = greenlet.getcurrent()
47
48    while level != 0:
49        if not isinstance(g, genlet):
50            raise RuntimeError('yield outside a genlet')
51        if level > 1:
52            g.parent.set_child(g)
53        g = g.parent
54        level -= 1
55
56    g.switch(value)
57
58
59def Genlet(func):
60    class TheGenlet(genlet):
61        fn = (func,)
62    return TheGenlet
63
64# ____________________________________________________________
65
66
67def g1(n, seen):
68    for i in range(n):
69        seen.append(i + 1)
70        yield i
71
72
73def g2(n, seen):
74    for i in range(n):
75        seen.append(i + 1)
76        Yield(i)
77
78g2 = Genlet(g2)
79
80
81def nested(i):
82    Yield(i)
83
84
85def g3(n, seen):
86    for i in range(n):
87        seen.append(i + 1)
88        nested(i)
89g3 = Genlet(g3)
90
91
92def a(n):
93    if n == 0:
94        return
95    for ii in ax(n - 1):
96        Yield(ii)
97    Yield(n)
98ax = Genlet(a)
99
100
101def perms(l):
102    if len(l) > 1:
103        for e in l:
104            # No syntactical sugar for generator expressions
105            x = [Yield([e] + p) for p in perms([x for x in l if x != e])]
106            assert x
107    else:
108        Yield(l)
109perms = Genlet(perms)
110
111
112def gr1(n):
113    for ii in range(1, n):
114        Yield(ii)
115        Yield(ii * ii, 2)
116
117gr1 = Genlet(gr1)
118
119
120def gr2(n, seen):
121    for ii in gr1(n):
122        seen.append(ii)
123
124gr2 = Genlet(gr2)
125
126
127class NestedGeneratorTests(TestCase):
128    def test_layered_genlets(self):
129        seen = []
130        for ii in gr2(5, seen):
131            seen.append(ii)
132        self.assertEqual(seen, [1, 1, 2, 4, 3, 9, 4, 16])
133
134    @fails_leakcheck
135    def test_permutations(self):
136        gen_perms = perms(list(range(4)))
137        permutations = list(gen_perms)
138        self.assertEqual(len(permutations), 4 * 3 * 2 * 1)
139        self.assertIn([0, 1, 2, 3], permutations)
140        self.assertIn([3, 2, 1, 0], permutations)
141        res = []
142        for ii in zip(perms(list(range(4))), perms(list(range(3)))):
143            res.append(ii)
144        self.assertEqual(
145            res,
146            [([0, 1, 2, 3], [0, 1, 2]), ([0, 1, 3, 2], [0, 2, 1]),
147             ([0, 2, 1, 3], [1, 0, 2]), ([0, 2, 3, 1], [1, 2, 0]),
148             ([0, 3, 1, 2], [2, 0, 1]), ([0, 3, 2, 1], [2, 1, 0])])
149        # XXX Test to make sure we are working as a generator expression
150
151    def test_genlet_simple(self):
152        for g in g1, g2, g3:
153            seen = []
154            for _ in range(3):
155                for j in g(5, seen):
156                    seen.append(j)
157            self.assertEqual(seen, 3 * [1, 0, 2, 1, 3, 2, 4, 3, 5, 4])
158
159    def test_genlet_bad(self):
160        try:
161            Yield(10)
162        except RuntimeError:
163            pass
164
165    def test_nested_genlets(self):
166        seen = []
167        for ii in ax(5):
168            seen.append(ii)
169