1# mode: run
2# tag: generators, gh3265
3
4try:
5    import backports_abc
6except ImportError: pass
7else: backports_abc.patch()
8
9try:
10    from collections.abc import Generator
11except ImportError:
12    try:
13        from collections import Generator
14    except ImportError:
15        Generator = object  # easy win
16
17
18def very_simple():
19    """
20    >>> x = very_simple()
21    >>> next(x)
22    1
23    >>> next(x)
24    Traceback (most recent call last):
25    StopIteration
26    >>> next(x)
27    Traceback (most recent call last):
28    StopIteration
29
30    >>> x = very_simple()
31    >>> x.send(1)
32    Traceback (most recent call last):
33    TypeError: can't send non-None value to a just-started generator
34    """
35    yield 1
36
37
38def attributes():
39    """
40    >>> x = attributes()
41    >>> x.__name__
42    'attributes'
43    >>> x.__qualname__
44    'attributes'
45    >>> x.gi_running  # before next()
46    False
47    >>> inner = next(x)
48    >>> x.gi_running  # after next()
49    False
50    >>> next(x)
51    Traceback (most recent call last):
52    StopIteration
53    >>> x.gi_running  # after termination
54    False
55
56    >>> y = inner()
57    >>> y.__name__
58    '<lambda>'
59    >>> y.__qualname__
60    'attributes.<locals>.inner.<locals>.<lambda>'
61
62    >>> y.__name__ = 123
63    Traceback (most recent call last):
64    TypeError: __name__ must be set to a string object
65    >>> y.__name__
66    '<lambda>'
67    >>> y.__qualname__ = None
68    Traceback (most recent call last):
69    TypeError: __qualname__ must be set to a string object
70    >>> y.__qualname__
71    'attributes.<locals>.inner.<locals>.<lambda>'
72
73    >>> y.__name__ = 'abc'
74    >>> y.__name__
75    'abc'
76    >>> y.__name__ = None
77    Traceback (most recent call last):
78    TypeError: __name__ must be set to a string object
79    >>> y.__name__
80    'abc'
81    >>> y.__qualname__ = 'huhu'
82    >>> y.__qualname__
83    'huhu'
84    >>> y.__qualname__ = 123
85    Traceback (most recent call last):
86    TypeError: __qualname__ must be set to a string object
87    >>> y.__qualname__
88    'huhu'
89    """
90    def inner():
91        return (lambda : (yield 1))
92    yield inner()
93
94
95def simple():
96    """
97    >>> x = simple()
98    >>> list(x)
99    [1, 2, 3]
100    """
101    yield 1
102    yield 2
103    yield 3
104
105def simple_seq(seq):
106    """
107    >>> x = simple_seq("abc")
108    >>> list(x)
109    ['a', 'b', 'c']
110    """
111    for i in seq:
112        yield i
113
114def simple_send():
115    """
116    >>> x = simple_send()
117    >>> next(x)
118    >>> x.send(1)
119    1
120    >>> x.send(2)
121    2
122    >>> x.send(3)
123    3
124    """
125    i = None
126    while True:
127        i = yield i
128
129def raising():
130    """
131    >>> x = raising()
132    >>> next(x)
133    Traceback (most recent call last):
134    KeyError: 'foo'
135    >>> next(x)
136    Traceback (most recent call last):
137    StopIteration
138    """
139    yield {}['foo']
140
141def with_outer(*args):
142    """
143    >>> x = with_outer(1, 2, 3)
144    >>> list(x())
145    [1, 2, 3]
146    """
147    def generator():
148        for i in args:
149            yield i
150    return generator
151
152def with_outer_raising(*args):
153    """
154    >>> x = with_outer_raising(1, 2, 3)
155    >>> list(x())
156    [1, 2, 3]
157    """
158    def generator():
159        for i in args:
160            yield i
161        raise StopIteration
162    return generator
163
164def test_close():
165    """
166    >>> x = test_close()
167    >>> x.close()
168    >>> x = test_close()
169    >>> next(x)
170    >>> x.close()
171    >>> next(x)
172    Traceback (most recent call last):
173    StopIteration
174    """
175    while True:
176        yield
177
178def test_ignore_close():
179    """
180    >>> x = test_ignore_close()
181    >>> x.close()
182    >>> x = test_ignore_close()
183    >>> next(x)
184    >>> x.close()
185    Traceback (most recent call last):
186    RuntimeError: generator ignored GeneratorExit
187    """
188    try:
189        yield
190    except GeneratorExit:
191        yield
192
193def check_throw():
194    """
195    >>> x = check_throw()
196    >>> x.throw(ValueError)
197    Traceback (most recent call last):
198    ValueError
199    >>> next(x)
200    Traceback (most recent call last):
201    StopIteration
202    >>> x = check_throw()
203    >>> next(x)
204    >>> x.throw(ValueError)
205    >>> next(x)
206    >>> x.throw(IndexError, "oops")
207    Traceback (most recent call last):
208    IndexError: oops
209    >>> next(x)
210    Traceback (most recent call last):
211    StopIteration
212    """
213    while True:
214        try:
215            yield
216        except ValueError:
217            pass
218
219def test_first_assignment():
220    """
221    >>> gen = test_first_assignment()
222    >>> next(gen)
223    5
224    >>> next(gen)
225    10
226    >>> next(gen)
227    (5, 10)
228    """
229    cdef x = 5 # first
230    yield x
231    cdef y = 10 # first
232    yield y
233    yield (x,y)
234
235def test_swap_assignment():
236    """
237    >>> gen = test_swap_assignment()
238    >>> next(gen)
239    (5, 10)
240    >>> next(gen)
241    (10, 5)
242    """
243    x,y = 5,10
244    yield (x,y)
245    x,y = y,x   # no ref-counting here
246    yield (x,y)
247
248
249class Foo(object):
250    """
251    >>> obj = Foo()
252    >>> list(obj.simple(1, 2, 3))
253    [1, 2, 3]
254    """
255    def simple(self, *args):
256        for i in args:
257            yield i
258
259def generator_nonlocal():
260    """
261    >>> g = generator_nonlocal()
262    >>> list(g(5))
263    [2, 3, 4, 5, 6]
264    """
265    def f(x):
266        def g(y):
267            nonlocal x
268            for i in range(y):
269                x += 1
270                yield x
271        return g
272    return f(1)
273
274def test_nested(a, b, c):
275    """
276    >>> obj = test_nested(1, 2, 3)
277    >>> [i() for i in obj]
278    [1, 2, 3, 4]
279    """
280    def one():
281        return a
282    def two():
283        return b
284    def three():
285        return c
286    def new_closure(a, b):
287        def sum():
288            return a + b
289        return sum
290    yield one
291    yield two
292    yield three
293    yield new_closure(a, c)
294
295
296def tolist(func):
297    def wrapper(*args, **kwargs):
298        return list(func(*args, **kwargs))
299    return wrapper
300
301@tolist
302def test_decorated(*args):
303    """
304    >>> test_decorated(1, 2, 3)
305    [1, 2, 3]
306    """
307    for i in args:
308        yield i
309
310
311def test_return(a):
312    """
313    >>> d = dict()
314    >>> obj = test_return(d)
315    >>> next(obj)
316    1
317    >>> next(obj)
318    Traceback (most recent call last):
319    StopIteration
320    >>> d['i_was_here']
321    True
322    """
323    yield 1
324    a['i_was_here'] = True
325    return
326
327
328def test_return_in_finally(a):
329    """
330    >>> d = dict()
331    >>> obj = test_return_in_finally(d)
332    >>> next(obj)
333    1
334    >>> next(obj)
335    Traceback (most recent call last):
336    StopIteration
337    >>> d['i_was_here']
338    True
339
340    >>> d = dict()
341    >>> obj = test_return_in_finally(d)
342    >>> next(obj)
343    1
344    >>> obj.send(2)
345    Traceback (most recent call last):
346    StopIteration
347    >>> d['i_was_here']
348    True
349
350    >>> obj = test_return_in_finally(None)
351    >>> next(obj)
352    1
353    >>> next(obj)
354    Traceback (most recent call last):
355    StopIteration
356
357    >>> obj = test_return_in_finally(None)
358    >>> next(obj)
359    1
360    >>> obj.send(2)
361    Traceback (most recent call last):
362    StopIteration
363    """
364    yield 1
365    try:
366        a['i_was_here'] = True
367    finally:
368        return
369
370
371def test_return_none_in_finally(a):
372    """
373    >>> d = dict()
374    >>> obj = test_return_none_in_finally(d)
375    >>> next(obj)
376    1
377    >>> next(obj)
378    Traceback (most recent call last):
379    StopIteration
380    >>> d['i_was_here']
381    True
382
383    >>> obj = test_return_none_in_finally(None)
384    >>> next(obj)
385    1
386    >>> next(obj)
387    Traceback (most recent call last):
388    StopIteration
389    """
390    yield 1
391    try:
392        a['i_was_here'] = True
393    finally:
394        return None
395
396
397def test_copied_yield(foo):
398    """
399    >>> class Manager(object):
400    ...    def __enter__(self):
401    ...        return self
402    ...    def __exit__(self, type, value, tb):
403    ...        pass
404    >>> list(test_copied_yield(Manager()))
405    [1]
406    """
407    with foo:
408        yield 1
409
410def test_nested_yield():
411    """
412    >>> obj = test_nested_yield()
413    >>> next(obj)
414    1
415    >>> obj.send(2)
416    2
417    >>> obj.send(3)
418    3
419    >>> obj.send(4)
420    Traceback (most recent call last):
421    StopIteration
422    """
423    yield (yield (yield 1))
424
425def test_inside_lambda():
426    """
427    >>> obj = test_inside_lambda()()
428    >>> next(obj)
429    1
430    >>> next(obj)
431    2
432    >>> next(obj)
433    Traceback (most recent call last):
434    StopIteration
435    """
436    return lambda:((yield 1), (yield 2))
437
438def test_nested_gen(int n):
439    """
440    >>> [list(a) for a in test_nested_gen(5)]
441    [[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]]
442    """
443    for a in range(n):
444        yield (b for b in range(a))
445
446def test_lambda(n):
447    """
448    >>> [i() for i in test_lambda(3)]
449    [0, 1, 2]
450    """
451    for i in range(n):
452        yield lambda : i
453
454
455def test_with_gil_section():
456    """
457    >>> list(test_with_gil_section())
458    [0, 1, 2]
459    """
460    cdef int i
461    with nogil:
462        for i in range(3):
463            with gil:
464                yield i
465
466
467def test_double_with_gil_section():
468    """
469    >>> list(test_double_with_gil_section())
470    [0, 1, 2, 3]
471    """
472    cdef int i,j
473    with nogil:
474        for i in range(2):
475            with gil:
476                with nogil:
477                    for j in range(2):
478                        with gil:
479                            yield i*2+j
480                with nogil:
481                    pass
482            with gil:
483                pass
484
485
486def test_generator_abc():
487    """
488    >>> isinstance(test_generator_abc(), Generator)
489    True
490
491    >>> try:
492    ...     from collections.abc import Generator
493    ... except ImportError:
494    ...     try:
495    ...         from collections import Generator
496    ...     except ImportError:
497    ...         Generator = object  # easy win
498
499    >>> isinstance(test_generator_abc(), Generator)
500    True
501    >>> isinstance((lambda:(yield))(), Generator)
502    True
503    """
504    yield 1
505
506
507def test_generator_frame(a=1):
508    """
509    >>> gen = test_generator_frame()
510    >>> import types
511    >>> isinstance(gen.gi_frame, types.FrameType) or gen.gi_frame
512    True
513    >>> gen.gi_frame is gen.gi_frame  # assert that it's cached
514    True
515    >>> gen.gi_frame.f_code is not None
516    True
517    >>> code_obj = gen.gi_frame.f_code
518    >>> code_obj.co_argcount
519    1
520    >>> code_obj.co_varnames
521    ('a', 'b')
522    """
523    b = a + 1
524    yield b
525
526
527# GH Issue 3265 - **kwds could cause a crash in some cases due to not
528# handling NULL pointers (in testing it shows as a REFNANNY error).
529# This was on creation of the generator and
530# doesn't really require it to be iterated through:
531
532def some_function():
533    return 0
534
535
536def test_generator_kwds1(**kwargs):
537    """
538    >>> for a in test_generator_kwds1():
539    ...     print(a)
540    0
541    """
542    yield some_function(**kwargs)
543
544
545def test_generator_kwds2(**kwargs):
546    """
547    >>> for a in test_generator_kwds2():
548    ...     print(a)
549    0
550    """
551    yield 0
552
553
554def test_generator_kwds3(**kwargs):
555    """
556    This didn't actually crash before but is still worth a try
557    >>> len(list(test_generator_kwds3()))
558    0
559    >>> for a in test_generator_kwds3(a=1):
560    ...    print(a)
561    a
562    """
563    yield from kwargs.keys()
564
565
566def test_generator_frame(a=1):
567    """
568    >>> gen = test_generator_frame()
569    >>> import types
570    >>> isinstance(gen.gi_frame, types.FrameType) or gen.gi_frame
571    True
572    >>> gen.gi_frame is gen.gi_frame  # assert that it's cached
573    True
574    >>> gen.gi_frame.f_code is not None
575    True
576    >>> code_obj = gen.gi_frame.f_code
577    >>> code_obj.co_argcount
578    1
579    >>> code_obj.co_varnames
580    ('a', 'b')
581    """
582    b = a + 1
583    yield b
584