1# mode: run
2# tag: generators
3
4import sys
5import cython
6
7
8def very_simple():
9    """
10    >>> x = very_simple()
11    >>> next(x)
12    1
13    >>> next(x)
14    Traceback (most recent call last):
15    StopIteration
16    >>> next(x)
17    Traceback (most recent call last):
18    StopIteration
19    >>> x = very_simple()
20    >>> x.send(1)
21    Traceback (most recent call last):
22    TypeError: can't send non-None value to a just-started generator
23    """
24    yield 1
25
26
27def simple():
28    """
29    >>> x = simple()
30    >>> list(x)
31    [1, 2, 3]
32    """
33    yield 1
34    yield 2
35    yield 3
36
37def simple_seq(seq):
38    """
39    >>> x = simple_seq("abc")
40    >>> list(x)
41    ['a', 'b', 'c']
42    """
43    for i in seq:
44        yield i
45
46def simple_send():
47    """
48    >>> x = simple_send()
49    >>> next(x)
50    >>> x.send(1)
51    1
52    >>> x.send(2)
53    2
54    >>> x.send(3)
55    3
56    """
57    i = None
58    while True:
59        i = yield i
60
61def raising():
62    """
63    >>> x = raising()
64    >>> next(x)
65    Traceback (most recent call last):
66    KeyError: 'foo'
67    >>> next(x)
68    Traceback (most recent call last):
69    StopIteration
70    """
71    yield {}['foo']
72
73def with_outer(*args):
74    """
75    >>> x = with_outer(1, 2, 3)
76    >>> list(x())
77    [1, 2, 3]
78    """
79    def generator():
80        for i in args:
81            yield i
82    return generator
83
84
85def test_close():
86    """
87    >>> x = test_close()
88    >>> x.close()
89    >>> x = test_close()
90    >>> next(x)
91    >>> x.close()
92    >>> next(x)
93    Traceback (most recent call last):
94    StopIteration
95    """
96    while True:
97        yield
98
99def test_ignore_close():
100    """
101    >>> x = test_ignore_close()
102    >>> x.close()
103    >>> x = test_ignore_close()
104    >>> next(x)
105    >>> x.close()
106    Traceback (most recent call last):
107    RuntimeError: generator ignored GeneratorExit
108    """
109    try:
110        yield
111    except GeneratorExit:
112        yield
113
114def check_throw():
115    """
116    >>> x = check_throw()
117    >>> x.throw(ValueError)
118    Traceback (most recent call last):
119    ValueError
120    >>> next(x)
121    Traceback (most recent call last):
122    StopIteration
123    >>> x = check_throw()
124    >>> next(x)
125    >>> x.throw(ValueError)
126    >>> next(x)
127    >>> x.throw(IndexError, "oops")
128    Traceback (most recent call last):
129    IndexError: oops
130    >>> next(x)
131    Traceback (most recent call last):
132    StopIteration
133    """
134    while True:
135        try:
136            yield
137        except ValueError:
138            pass
139
140
141def check_yield_in_except():
142    """
143    >>> if sys.version_info[0] == 2: sys.exc_clear()
144    >>> try:
145    ...     raise TypeError("RAISED !")
146    ... except TypeError as orig_exc:
147    ...     assert isinstance(orig_exc, TypeError), orig_exc
148    ...     g = check_yield_in_except()
149    ...     print(orig_exc is sys.exc_info()[1] or sys.exc_info())
150    ...     next(g)
151    ...     print(orig_exc is sys.exc_info()[1] or sys.exc_info())
152    ...     next(g)
153    ...     print(orig_exc is sys.exc_info()[1] or sys.exc_info())
154    True
155    True
156    True
157    >>> next(g)
158    Traceback (most recent call last):
159    StopIteration
160    """
161    try:
162        yield
163        raise ValueError
164    except ValueError as exc:
165        assert sys.exc_info()[1] is exc, sys.exc_info()
166        yield
167        if cython.compiled or sys.version_info[0] > 2:
168            assert sys.exc_info()[1] is exc, sys.exc_info()
169
170
171def yield_in_except_throw_exc_type():
172    """
173    >>> g = yield_in_except_throw_exc_type()
174    >>> next(g)
175    >>> g.throw(TypeError)
176    Traceback (most recent call last):
177    TypeError
178    >>> next(g)
179    Traceback (most recent call last):
180    StopIteration
181    """
182    try:
183        raise ValueError
184    except ValueError as exc:
185        assert sys.exc_info()[1] is exc, sys.exc_info()
186        yield
187        assert sys.exc_info()[1] is exc, sys.exc_info()
188
189
190def yield_in_except_throw_instance():
191    """
192    >>> g = yield_in_except_throw_instance()
193    >>> next(g)
194    >>> g.throw(TypeError())
195    Traceback (most recent call last):
196    TypeError
197    >>> next(g)
198    Traceback (most recent call last):
199    StopIteration
200    """
201    try:
202        raise ValueError
203    except ValueError as exc:
204        assert sys.exc_info()[1] is exc, sys.exc_info()
205        yield
206        assert sys.exc_info()[1] is exc, sys.exc_info()
207
208
209def test_swap_assignment():
210    """
211    >>> gen = test_swap_assignment()
212    >>> next(gen)
213    (5, 10)
214    >>> next(gen)
215    (10, 5)
216    """
217    x,y = 5,10
218    yield (x,y)
219    x,y = y,x   # no ref-counting here
220    yield (x,y)
221
222
223class Foo(object):
224    """
225    >>> obj = Foo()
226    >>> list(obj.simple(1, 2, 3))
227    [1, 2, 3]
228    """
229    def simple(self, *args):
230        for i in args:
231            yield i
232
233def test_nested(a, b, c):
234    """
235    >>> obj = test_nested(1, 2, 3)
236    >>> [i() for i in obj]
237    [1, 2, 3, 4]
238    """
239    def one():
240        return a
241    def two():
242        return b
243    def three():
244        return c
245    def new_closure(a, b):
246        def sum():
247            return a + b
248        return sum
249    yield one
250    yield two
251    yield three
252    yield new_closure(a, c)
253
254
255def tolist(func):
256    def wrapper(*args, **kwargs):
257        return list(func(*args, **kwargs))
258    return wrapper
259
260@tolist
261def test_decorated(*args):
262    """
263    >>> test_decorated(1, 2, 3)
264    [1, 2, 3]
265    """
266    for i in args:
267        yield i
268
269def test_return(a):
270    """
271    >>> d = dict()
272    >>> obj = test_return(d)
273    >>> next(obj)
274    1
275    >>> next(obj)
276    Traceback (most recent call last):
277    StopIteration
278    >>> d['i_was_here']
279    True
280    """
281    yield 1
282    a['i_was_here'] = True
283    return
284
285def test_copied_yield(foo):
286    """
287    >>> class Manager(object):
288    ...    def __enter__(self):
289    ...        return self
290    ...    def __exit__(self, type, value, tb):
291    ...        pass
292    >>> list(test_copied_yield(Manager()))
293    [1]
294    """
295    with foo:
296        yield 1
297
298def test_nested_yield():
299    """
300    >>> obj = test_nested_yield()
301    >>> next(obj)
302    1
303    >>> obj.send(2)
304    2
305    >>> obj.send(3)
306    3
307    >>> obj.send(4)
308    Traceback (most recent call last):
309    StopIteration
310    """
311    yield (yield (yield 1))
312
313def test_sum_of_yields(n):
314    """
315    >>> g = test_sum_of_yields(3)
316    >>> next(g)
317    (0, 0)
318    >>> g.send(1)
319    (0, 1)
320    >>> g.send(1)
321    (1, 2)
322    """
323    x = 0
324    x += yield (0, x)
325    x += yield (0, x)
326    yield (1, x)
327
328def test_nested_gen(n):
329    """
330    >>> [list(a) for a in test_nested_gen(5)]
331    [[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]]
332    """
333    for a in range(n):
334        yield (b for b in range(a))
335
336def test_lambda(n):
337    """
338    >>> [i() for i in test_lambda(3)]
339    [0, 1, 2]
340    """
341    for i in range(n):
342        yield lambda : i
343
344def test_generator_cleanup():
345    """
346    >>> g = test_generator_cleanup()
347    >>> del g
348    >>> g = test_generator_cleanup()
349    >>> next(g)
350    1
351    >>> del g
352    cleanup
353    """
354    try:
355        yield 1
356    finally:
357        print('cleanup')
358
359def test_del_in_generator():
360    """
361    >>> [ s for s in test_del_in_generator() ]
362    ['abcabcabc', 'abcabcabc']
363    """
364    x = len('abc') * 'abc'
365    a = x
366    yield x
367    del x
368    yield a
369    del a
370
371@cython.test_fail_if_path_exists("//IfStatNode", "//PrintStatNode")
372def test_yield_in_const_conditional_false():
373    """
374    >>> list(test_yield_in_const_conditional_false())
375    []
376    """
377    if False:
378        print((yield 1))
379
380@cython.test_fail_if_path_exists("//IfStatNode")
381@cython.test_assert_path_exists("//PrintStatNode")
382def test_yield_in_const_conditional_true():
383    """
384    >>> list(test_yield_in_const_conditional_true())
385    None
386    [1]
387    """
388    if True:
389        print((yield 1))
390