1import inspect
2import types
3import unittest
4
5from test.support import import_module
6asyncio = import_module("asyncio")
7
8
9class AwaitException(Exception):
10    pass
11
12
13@types.coroutine
14def awaitable(*, throw=False):
15    if throw:
16        yield ('throw',)
17    else:
18        yield ('result',)
19
20
21def run_until_complete(coro):
22    exc = False
23    while True:
24        try:
25            if exc:
26                exc = False
27                fut = coro.throw(AwaitException)
28            else:
29                fut = coro.send(None)
30        except StopIteration as ex:
31            return ex.args[0]
32
33        if fut == ('throw',):
34            exc = True
35
36
37def to_list(gen):
38    async def iterate():
39        res = []
40        async for i in gen:
41            res.append(i)
42        return res
43
44    return run_until_complete(iterate())
45
46
47class AsyncGenSyntaxTest(unittest.TestCase):
48
49    def test_async_gen_syntax_01(self):
50        code = '''async def foo():
51            await abc
52            yield from 123
53        '''
54
55        with self.assertRaisesRegex(SyntaxError, 'yield from.*inside async'):
56            exec(code, {}, {})
57
58    def test_async_gen_syntax_02(self):
59        code = '''async def foo():
60            yield from 123
61        '''
62
63        with self.assertRaisesRegex(SyntaxError, 'yield from.*inside async'):
64            exec(code, {}, {})
65
66    def test_async_gen_syntax_03(self):
67        code = '''async def foo():
68            await abc
69            yield
70            return 123
71        '''
72
73        with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'):
74            exec(code, {}, {})
75
76    def test_async_gen_syntax_04(self):
77        code = '''async def foo():
78            yield
79            return 123
80        '''
81
82        with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'):
83            exec(code, {}, {})
84
85    def test_async_gen_syntax_05(self):
86        code = '''async def foo():
87            if 0:
88                yield
89            return 12
90        '''
91
92        with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'):
93            exec(code, {}, {})
94
95
96class AsyncGenTest(unittest.TestCase):
97
98    def compare_generators(self, sync_gen, async_gen):
99        def sync_iterate(g):
100            res = []
101            while True:
102                try:
103                    res.append(g.__next__())
104                except StopIteration:
105                    res.append('STOP')
106                    break
107                except Exception as ex:
108                    res.append(str(type(ex)))
109            return res
110
111        def async_iterate(g):
112            res = []
113            while True:
114                an = g.__anext__()
115                try:
116                    while True:
117                        try:
118                            an.__next__()
119                        except StopIteration as ex:
120                            if ex.args:
121                                res.append(ex.args[0])
122                                break
123                            else:
124                                res.append('EMPTY StopIteration')
125                                break
126                        except StopAsyncIteration:
127                            raise
128                        except Exception as ex:
129                            res.append(str(type(ex)))
130                            break
131                except StopAsyncIteration:
132                    res.append('STOP')
133                    break
134            return res
135
136        sync_gen_result = sync_iterate(sync_gen)
137        async_gen_result = async_iterate(async_gen)
138        self.assertEqual(sync_gen_result, async_gen_result)
139        return async_gen_result
140
141    def test_async_gen_iteration_01(self):
142        async def gen():
143            await awaitable()
144            a = yield 123
145            self.assertIs(a, None)
146            await awaitable()
147            yield 456
148            await awaitable()
149            yield 789
150
151        self.assertEqual(to_list(gen()), [123, 456, 789])
152
153    def test_async_gen_iteration_02(self):
154        async def gen():
155            await awaitable()
156            yield 123
157            await awaitable()
158
159        g = gen()
160        ai = g.__aiter__()
161
162        an = ai.__anext__()
163        self.assertEqual(an.__next__(), ('result',))
164
165        try:
166            an.__next__()
167        except StopIteration as ex:
168            self.assertEqual(ex.args[0], 123)
169        else:
170            self.fail('StopIteration was not raised')
171
172        an = ai.__anext__()
173        self.assertEqual(an.__next__(), ('result',))
174
175        try:
176            an.__next__()
177        except StopAsyncIteration as ex:
178            self.assertFalse(ex.args)
179        else:
180            self.fail('StopAsyncIteration was not raised')
181
182    def test_async_gen_exception_03(self):
183        async def gen():
184            await awaitable()
185            yield 123
186            await awaitable(throw=True)
187            yield 456
188
189        with self.assertRaises(AwaitException):
190            to_list(gen())
191
192    def test_async_gen_exception_04(self):
193        async def gen():
194            await awaitable()
195            yield 123
196            1 / 0
197
198        g = gen()
199        ai = g.__aiter__()
200        an = ai.__anext__()
201        self.assertEqual(an.__next__(), ('result',))
202
203        try:
204            an.__next__()
205        except StopIteration as ex:
206            self.assertEqual(ex.args[0], 123)
207        else:
208            self.fail('StopIteration was not raised')
209
210        with self.assertRaises(ZeroDivisionError):
211            ai.__anext__().__next__()
212
213    def test_async_gen_exception_05(self):
214        async def gen():
215            yield 123
216            raise StopAsyncIteration
217
218        with self.assertRaisesRegex(RuntimeError,
219                                    'async generator.*StopAsyncIteration'):
220            to_list(gen())
221
222    def test_async_gen_exception_06(self):
223        async def gen():
224            yield 123
225            raise StopIteration
226
227        with self.assertRaisesRegex(RuntimeError,
228                                    'async generator.*StopIteration'):
229            to_list(gen())
230
231    def test_async_gen_exception_07(self):
232        def sync_gen():
233            try:
234                yield 1
235                1 / 0
236            finally:
237                yield 2
238                yield 3
239
240            yield 100
241
242        async def async_gen():
243            try:
244                yield 1
245                1 / 0
246            finally:
247                yield 2
248                yield 3
249
250            yield 100
251
252        self.compare_generators(sync_gen(), async_gen())
253
254    def test_async_gen_exception_08(self):
255        def sync_gen():
256            try:
257                yield 1
258            finally:
259                yield 2
260                1 / 0
261                yield 3
262
263            yield 100
264
265        async def async_gen():
266            try:
267                yield 1
268                await awaitable()
269            finally:
270                await awaitable()
271                yield 2
272                1 / 0
273                yield 3
274
275            yield 100
276
277        self.compare_generators(sync_gen(), async_gen())
278
279    def test_async_gen_exception_09(self):
280        def sync_gen():
281            try:
282                yield 1
283                1 / 0
284            finally:
285                yield 2
286                yield 3
287
288            yield 100
289
290        async def async_gen():
291            try:
292                await awaitable()
293                yield 1
294                1 / 0
295            finally:
296                yield 2
297                await awaitable()
298                yield 3
299
300            yield 100
301
302        self.compare_generators(sync_gen(), async_gen())
303
304    def test_async_gen_exception_10(self):
305        async def gen():
306            yield 123
307        with self.assertRaisesRegex(TypeError,
308                                    "non-None value .* async generator"):
309            gen().__anext__().send(100)
310
311    def test_async_gen_exception_11(self):
312        def sync_gen():
313            yield 10
314            yield 20
315
316        def sync_gen_wrapper():
317            yield 1
318            sg = sync_gen()
319            sg.send(None)
320            try:
321                sg.throw(GeneratorExit())
322            except GeneratorExit:
323                yield 2
324            yield 3
325
326        async def async_gen():
327            yield 10
328            yield 20
329
330        async def async_gen_wrapper():
331            yield 1
332            asg = async_gen()
333            await asg.asend(None)
334            try:
335                await asg.athrow(GeneratorExit())
336            except GeneratorExit:
337                yield 2
338            yield 3
339
340        self.compare_generators(sync_gen_wrapper(), async_gen_wrapper())
341
342    def test_async_gen_api_01(self):
343        async def gen():
344            yield 123
345
346        g = gen()
347
348        self.assertEqual(g.__name__, 'gen')
349        g.__name__ = '123'
350        self.assertEqual(g.__name__, '123')
351
352        self.assertIn('.gen', g.__qualname__)
353        g.__qualname__ = '123'
354        self.assertEqual(g.__qualname__, '123')
355
356        self.assertIsNone(g.ag_await)
357        self.assertIsInstance(g.ag_frame, types.FrameType)
358        self.assertFalse(g.ag_running)
359        self.assertIsInstance(g.ag_code, types.CodeType)
360
361        self.assertTrue(inspect.isawaitable(g.aclose()))
362
363
364class AsyncGenAsyncioTest(unittest.TestCase):
365
366    def setUp(self):
367        self.loop = asyncio.new_event_loop()
368        asyncio.set_event_loop(None)
369
370    def tearDown(self):
371        self.loop.close()
372        self.loop = None
373        asyncio.set_event_loop_policy(None)
374
375    async def to_list(self, gen):
376        res = []
377        async for i in gen:
378            res.append(i)
379        return res
380
381    def test_async_gen_asyncio_01(self):
382        async def gen():
383            yield 1
384            await asyncio.sleep(0.01)
385            yield 2
386            await asyncio.sleep(0.01)
387            return
388            yield 3
389
390        res = self.loop.run_until_complete(self.to_list(gen()))
391        self.assertEqual(res, [1, 2])
392
393    def test_async_gen_asyncio_02(self):
394        async def gen():
395            yield 1
396            await asyncio.sleep(0.01)
397            yield 2
398            1 / 0
399            yield 3
400
401        with self.assertRaises(ZeroDivisionError):
402            self.loop.run_until_complete(self.to_list(gen()))
403
404    def test_async_gen_asyncio_03(self):
405        loop = self.loop
406
407        class Gen:
408            async def __aiter__(self):
409                yield 1
410                await asyncio.sleep(0.01)
411                yield 2
412
413        res = loop.run_until_complete(self.to_list(Gen()))
414        self.assertEqual(res, [1, 2])
415
416    def test_async_gen_asyncio_anext_04(self):
417        async def foo():
418            yield 1
419            await asyncio.sleep(0.01)
420            try:
421                yield 2
422                yield 3
423            except ZeroDivisionError:
424                yield 1000
425            await asyncio.sleep(0.01)
426            yield 4
427
428        async def run1():
429            it = foo().__aiter__()
430
431            self.assertEqual(await it.__anext__(), 1)
432            self.assertEqual(await it.__anext__(), 2)
433            self.assertEqual(await it.__anext__(), 3)
434            self.assertEqual(await it.__anext__(), 4)
435            with self.assertRaises(StopAsyncIteration):
436                await it.__anext__()
437            with self.assertRaises(StopAsyncIteration):
438                await it.__anext__()
439
440        async def run2():
441            it = foo().__aiter__()
442
443            self.assertEqual(await it.__anext__(), 1)
444            self.assertEqual(await it.__anext__(), 2)
445            try:
446                it.__anext__().throw(ZeroDivisionError)
447            except StopIteration as ex:
448                self.assertEqual(ex.args[0], 1000)
449            else:
450                self.fail('StopIteration was not raised')
451            self.assertEqual(await it.__anext__(), 4)
452            with self.assertRaises(StopAsyncIteration):
453                await it.__anext__()
454
455        self.loop.run_until_complete(run1())
456        self.loop.run_until_complete(run2())
457
458    def test_async_gen_asyncio_anext_05(self):
459        async def foo():
460            v = yield 1
461            v = yield v
462            yield v * 100
463
464        async def run():
465            it = foo().__aiter__()
466
467            try:
468                it.__anext__().send(None)
469            except StopIteration as ex:
470                self.assertEqual(ex.args[0], 1)
471            else:
472                self.fail('StopIteration was not raised')
473
474            try:
475                it.__anext__().send(10)
476            except StopIteration as ex:
477                self.assertEqual(ex.args[0], 10)
478            else:
479                self.fail('StopIteration was not raised')
480
481            try:
482                it.__anext__().send(12)
483            except StopIteration as ex:
484                self.assertEqual(ex.args[0], 1200)
485            else:
486                self.fail('StopIteration was not raised')
487
488            with self.assertRaises(StopAsyncIteration):
489                await it.__anext__()
490
491        self.loop.run_until_complete(run())
492
493    def test_async_gen_asyncio_anext_06(self):
494        DONE = 0
495
496        # test synchronous generators
497        def foo():
498            try:
499                yield
500            except:
501                pass
502        g = foo()
503        g.send(None)
504        with self.assertRaises(StopIteration):
505            g.send(None)
506
507        # now with asynchronous generators
508
509        async def gen():
510            nonlocal DONE
511            try:
512                yield
513            except:
514                pass
515            DONE = 1
516
517        async def run():
518            nonlocal DONE
519            g = gen()
520            await g.asend(None)
521            with self.assertRaises(StopAsyncIteration):
522                await g.asend(None)
523            DONE += 10
524
525        self.loop.run_until_complete(run())
526        self.assertEqual(DONE, 11)
527
528    def test_async_gen_asyncio_anext_tuple(self):
529        async def foo():
530            try:
531                yield (1,)
532            except ZeroDivisionError:
533                yield (2,)
534
535        async def run():
536            it = foo().__aiter__()
537
538            self.assertEqual(await it.__anext__(), (1,))
539            with self.assertRaises(StopIteration) as cm:
540                it.__anext__().throw(ZeroDivisionError)
541            self.assertEqual(cm.exception.args[0], (2,))
542            with self.assertRaises(StopAsyncIteration):
543                await it.__anext__()
544
545        self.loop.run_until_complete(run())
546
547    def test_async_gen_asyncio_anext_stopiteration(self):
548        async def foo():
549            try:
550                yield StopIteration(1)
551            except ZeroDivisionError:
552                yield StopIteration(3)
553
554        async def run():
555            it = foo().__aiter__()
556
557            v = await it.__anext__()
558            self.assertIsInstance(v, StopIteration)
559            self.assertEqual(v.value, 1)
560            with self.assertRaises(StopIteration) as cm:
561                it.__anext__().throw(ZeroDivisionError)
562            v = cm.exception.args[0]
563            self.assertIsInstance(v, StopIteration)
564            self.assertEqual(v.value, 3)
565            with self.assertRaises(StopAsyncIteration):
566                await it.__anext__()
567
568        self.loop.run_until_complete(run())
569
570    def test_async_gen_asyncio_aclose_06(self):
571        async def foo():
572            try:
573                yield 1
574                1 / 0
575            finally:
576                await asyncio.sleep(0.01)
577                yield 12
578
579        async def run():
580            gen = foo()
581            it = gen.__aiter__()
582            await it.__anext__()
583            await gen.aclose()
584
585        with self.assertRaisesRegex(
586                RuntimeError,
587                "async generator ignored GeneratorExit"):
588            self.loop.run_until_complete(run())
589
590    def test_async_gen_asyncio_aclose_07(self):
591        DONE = 0
592
593        async def foo():
594            nonlocal DONE
595            try:
596                yield 1
597                1 / 0
598            finally:
599                await asyncio.sleep(0.01)
600                await asyncio.sleep(0.01)
601                DONE += 1
602            DONE += 1000
603
604        async def run():
605            gen = foo()
606            it = gen.__aiter__()
607            await it.__anext__()
608            await gen.aclose()
609
610        self.loop.run_until_complete(run())
611        self.assertEqual(DONE, 1)
612
613    def test_async_gen_asyncio_aclose_08(self):
614        DONE = 0
615
616        fut = asyncio.Future(loop=self.loop)
617
618        async def foo():
619            nonlocal DONE
620            try:
621                yield 1
622                await fut
623                DONE += 1000
624                yield 2
625            finally:
626                await asyncio.sleep(0.01)
627                await asyncio.sleep(0.01)
628                DONE += 1
629            DONE += 1000
630
631        async def run():
632            gen = foo()
633            it = gen.__aiter__()
634            self.assertEqual(await it.__anext__(), 1)
635            await gen.aclose()
636
637        self.loop.run_until_complete(run())
638        self.assertEqual(DONE, 1)
639
640        # Silence ResourceWarnings
641        fut.cancel()
642        self.loop.run_until_complete(asyncio.sleep(0.01))
643
644    def test_async_gen_asyncio_gc_aclose_09(self):
645        DONE = 0
646
647        async def gen():
648            nonlocal DONE
649            try:
650                while True:
651                    yield 1
652            finally:
653                await asyncio.sleep(0.01)
654                await asyncio.sleep(0.01)
655                DONE = 1
656
657        async def run():
658            g = gen()
659            await g.__anext__()
660            await g.__anext__()
661            del g
662
663            await asyncio.sleep(0.1)
664
665        self.loop.run_until_complete(run())
666        self.assertEqual(DONE, 1)
667
668    def test_async_gen_asyncio_aclose_10(self):
669        DONE = 0
670
671        # test synchronous generators
672        def foo():
673            try:
674                yield
675            except:
676                pass
677        g = foo()
678        g.send(None)
679        g.close()
680
681        # now with asynchronous generators
682
683        async def gen():
684            nonlocal DONE
685            try:
686                yield
687            except:
688                pass
689            DONE = 1
690
691        async def run():
692            nonlocal DONE
693            g = gen()
694            await g.asend(None)
695            await g.aclose()
696            DONE += 10
697
698        self.loop.run_until_complete(run())
699        self.assertEqual(DONE, 11)
700
701    def test_async_gen_asyncio_aclose_11(self):
702        DONE = 0
703
704        # test synchronous generators
705        def foo():
706            try:
707                yield
708            except:
709                pass
710            yield
711        g = foo()
712        g.send(None)
713        with self.assertRaisesRegex(RuntimeError, 'ignored GeneratorExit'):
714            g.close()
715
716        # now with asynchronous generators
717
718        async def gen():
719            nonlocal DONE
720            try:
721                yield
722            except:
723                pass
724            yield
725            DONE += 1
726
727        async def run():
728            nonlocal DONE
729            g = gen()
730            await g.asend(None)
731            with self.assertRaisesRegex(RuntimeError, 'ignored GeneratorExit'):
732                await g.aclose()
733            DONE += 10
734
735        self.loop.run_until_complete(run())
736        self.assertEqual(DONE, 10)
737
738    def test_async_gen_asyncio_aclose_12(self):
739        DONE = 0
740
741        async def target():
742            await asyncio.sleep(0.01)
743            1 / 0
744
745        async def foo():
746            nonlocal DONE
747            task = asyncio.create_task(target())
748            try:
749                yield 1
750            finally:
751                try:
752                    await task
753                except ZeroDivisionError:
754                    DONE = 1
755
756        async def run():
757            gen = foo()
758            it = gen.__aiter__()
759            await it.__anext__()
760            await gen.aclose()
761
762        self.loop.run_until_complete(run())
763        self.assertEqual(DONE, 1)
764
765    def test_async_gen_asyncio_asend_01(self):
766        DONE = 0
767
768        # Sanity check:
769        def sgen():
770            v = yield 1
771            yield v * 2
772        sg = sgen()
773        v = sg.send(None)
774        self.assertEqual(v, 1)
775        v = sg.send(100)
776        self.assertEqual(v, 200)
777
778        async def gen():
779            nonlocal DONE
780            try:
781                await asyncio.sleep(0.01)
782                v = yield 1
783                await asyncio.sleep(0.01)
784                yield v * 2
785                await asyncio.sleep(0.01)
786                return
787            finally:
788                await asyncio.sleep(0.01)
789                await asyncio.sleep(0.01)
790                DONE = 1
791
792        async def run():
793            g = gen()
794
795            v = await g.asend(None)
796            self.assertEqual(v, 1)
797
798            v = await g.asend(100)
799            self.assertEqual(v, 200)
800
801            with self.assertRaises(StopAsyncIteration):
802                await g.asend(None)
803
804        self.loop.run_until_complete(run())
805        self.assertEqual(DONE, 1)
806
807    def test_async_gen_asyncio_asend_02(self):
808        DONE = 0
809
810        async def sleep_n_crash(delay):
811            await asyncio.sleep(delay)
812            1 / 0
813
814        async def gen():
815            nonlocal DONE
816            try:
817                await asyncio.sleep(0.01)
818                v = yield 1
819                await sleep_n_crash(0.01)
820                DONE += 1000
821                yield v * 2
822            finally:
823                await asyncio.sleep(0.01)
824                await asyncio.sleep(0.01)
825                DONE = 1
826
827        async def run():
828            g = gen()
829
830            v = await g.asend(None)
831            self.assertEqual(v, 1)
832
833            await g.asend(100)
834
835        with self.assertRaises(ZeroDivisionError):
836            self.loop.run_until_complete(run())
837        self.assertEqual(DONE, 1)
838
839    def test_async_gen_asyncio_asend_03(self):
840        DONE = 0
841
842        async def sleep_n_crash(delay):
843            fut = asyncio.ensure_future(asyncio.sleep(delay),
844                                        loop=self.loop)
845            self.loop.call_later(delay / 2, lambda: fut.cancel())
846            return await fut
847
848        async def gen():
849            nonlocal DONE
850            try:
851                await asyncio.sleep(0.01)
852                v = yield 1
853                await sleep_n_crash(0.01)
854                DONE += 1000
855                yield v * 2
856            finally:
857                await asyncio.sleep(0.01)
858                await asyncio.sleep(0.01)
859                DONE = 1
860
861        async def run():
862            g = gen()
863
864            v = await g.asend(None)
865            self.assertEqual(v, 1)
866
867            await g.asend(100)
868
869        with self.assertRaises(asyncio.CancelledError):
870            self.loop.run_until_complete(run())
871        self.assertEqual(DONE, 1)
872
873    def test_async_gen_asyncio_athrow_01(self):
874        DONE = 0
875
876        class FooEr(Exception):
877            pass
878
879        # Sanity check:
880        def sgen():
881            try:
882                v = yield 1
883            except FooEr:
884                v = 1000
885            yield v * 2
886        sg = sgen()
887        v = sg.send(None)
888        self.assertEqual(v, 1)
889        v = sg.throw(FooEr)
890        self.assertEqual(v, 2000)
891        with self.assertRaises(StopIteration):
892            sg.send(None)
893
894        async def gen():
895            nonlocal DONE
896            try:
897                await asyncio.sleep(0.01)
898                try:
899                    v = yield 1
900                except FooEr:
901                    v = 1000
902                    await asyncio.sleep(0.01)
903                yield v * 2
904                await asyncio.sleep(0.01)
905                # return
906            finally:
907                await asyncio.sleep(0.01)
908                await asyncio.sleep(0.01)
909                DONE = 1
910
911        async def run():
912            g = gen()
913
914            v = await g.asend(None)
915            self.assertEqual(v, 1)
916
917            v = await g.athrow(FooEr)
918            self.assertEqual(v, 2000)
919
920            with self.assertRaises(StopAsyncIteration):
921                await g.asend(None)
922
923        self.loop.run_until_complete(run())
924        self.assertEqual(DONE, 1)
925
926    def test_async_gen_asyncio_athrow_02(self):
927        DONE = 0
928
929        class FooEr(Exception):
930            pass
931
932        async def sleep_n_crash(delay):
933            fut = asyncio.ensure_future(asyncio.sleep(delay),
934                                        loop=self.loop)
935            self.loop.call_later(delay / 2, lambda: fut.cancel())
936            return await fut
937
938        async def gen():
939            nonlocal DONE
940            try:
941                await asyncio.sleep(0.01)
942                try:
943                    v = yield 1
944                except FooEr:
945                    await sleep_n_crash(0.01)
946                yield v * 2
947                await asyncio.sleep(0.01)
948                # return
949            finally:
950                await asyncio.sleep(0.01)
951                await asyncio.sleep(0.01)
952                DONE = 1
953
954        async def run():
955            g = gen()
956
957            v = await g.asend(None)
958            self.assertEqual(v, 1)
959
960            try:
961                await g.athrow(FooEr)
962            except asyncio.CancelledError:
963                self.assertEqual(DONE, 1)
964                raise
965            else:
966                self.fail('CancelledError was not raised')
967
968        with self.assertRaises(asyncio.CancelledError):
969            self.loop.run_until_complete(run())
970        self.assertEqual(DONE, 1)
971
972    def test_async_gen_asyncio_athrow_03(self):
973        DONE = 0
974
975        # test synchronous generators
976        def foo():
977            try:
978                yield
979            except:
980                pass
981        g = foo()
982        g.send(None)
983        with self.assertRaises(StopIteration):
984            g.throw(ValueError)
985
986        # now with asynchronous generators
987
988        async def gen():
989            nonlocal DONE
990            try:
991                yield
992            except:
993                pass
994            DONE = 1
995
996        async def run():
997            nonlocal DONE
998            g = gen()
999            await g.asend(None)
1000            with self.assertRaises(StopAsyncIteration):
1001                await g.athrow(ValueError)
1002            DONE += 10
1003
1004        self.loop.run_until_complete(run())
1005        self.assertEqual(DONE, 11)
1006
1007    def test_async_gen_asyncio_athrow_tuple(self):
1008        async def gen():
1009            try:
1010                yield 1
1011            except ZeroDivisionError:
1012                yield (2,)
1013
1014        async def run():
1015            g = gen()
1016            v = await g.asend(None)
1017            self.assertEqual(v, 1)
1018            v = await g.athrow(ZeroDivisionError)
1019            self.assertEqual(v, (2,))
1020            with self.assertRaises(StopAsyncIteration):
1021                await g.asend(None)
1022
1023        self.loop.run_until_complete(run())
1024
1025    def test_async_gen_asyncio_athrow_stopiteration(self):
1026        async def gen():
1027            try:
1028                yield 1
1029            except ZeroDivisionError:
1030                yield StopIteration(2)
1031
1032        async def run():
1033            g = gen()
1034            v = await g.asend(None)
1035            self.assertEqual(v, 1)
1036            v = await g.athrow(ZeroDivisionError)
1037            self.assertIsInstance(v, StopIteration)
1038            self.assertEqual(v.value, 2)
1039            with self.assertRaises(StopAsyncIteration):
1040                await g.asend(None)
1041
1042        self.loop.run_until_complete(run())
1043
1044    def test_async_gen_asyncio_shutdown_01(self):
1045        finalized = 0
1046
1047        async def waiter(timeout):
1048            nonlocal finalized
1049            try:
1050                await asyncio.sleep(timeout)
1051                yield 1
1052            finally:
1053                await asyncio.sleep(0)
1054                finalized += 1
1055
1056        async def wait():
1057            async for _ in waiter(1):
1058                pass
1059
1060        t1 = self.loop.create_task(wait())
1061        t2 = self.loop.create_task(wait())
1062
1063        self.loop.run_until_complete(asyncio.sleep(0.1))
1064
1065        # Silence warnings
1066        t1.cancel()
1067        t2.cancel()
1068
1069        with self.assertRaises(asyncio.CancelledError):
1070            self.loop.run_until_complete(t1)
1071        with self.assertRaises(asyncio.CancelledError):
1072            self.loop.run_until_complete(t2)
1073
1074        self.loop.run_until_complete(self.loop.shutdown_asyncgens())
1075
1076        self.assertEqual(finalized, 2)
1077
1078    def test_async_gen_expression_01(self):
1079        async def arange(n):
1080            for i in range(n):
1081                await asyncio.sleep(0.01)
1082                yield i
1083
1084        def make_arange(n):
1085            # This syntax is legal starting with Python 3.7
1086            return (i * 2 async for i in arange(n))
1087
1088        async def run():
1089            return [i async for i in make_arange(10)]
1090
1091        res = self.loop.run_until_complete(run())
1092        self.assertEqual(res, [i * 2 for i in range(10)])
1093
1094    def test_async_gen_expression_02(self):
1095        async def wrap(n):
1096            await asyncio.sleep(0.01)
1097            return n
1098
1099        def make_arange(n):
1100            # This syntax is legal starting with Python 3.7
1101            return (i * 2 for i in range(n) if await wrap(i))
1102
1103        async def run():
1104            return [i async for i in make_arange(10)]
1105
1106        res = self.loop.run_until_complete(run())
1107        self.assertEqual(res, [i * 2 for i in range(1, 10)])
1108
1109    def test_asyncgen_nonstarted_hooks_are_cancellable(self):
1110        # See https://bugs.python.org/issue38013
1111        messages = []
1112
1113        def exception_handler(loop, context):
1114            messages.append(context)
1115
1116        async def async_iterate():
1117            yield 1
1118            yield 2
1119
1120        async def main():
1121            loop = asyncio.get_running_loop()
1122            loop.set_exception_handler(exception_handler)
1123
1124            async for i in async_iterate():
1125                break
1126
1127        asyncio.run(main())
1128
1129        self.assertEqual([], messages)
1130
1131    def test_async_gen_await_same_anext_coro_twice(self):
1132        async def async_iterate():
1133            yield 1
1134            yield 2
1135
1136        async def run():
1137            it = async_iterate()
1138            nxt = it.__anext__()
1139            await nxt
1140            with self.assertRaisesRegex(
1141                    RuntimeError,
1142                    r"cannot reuse already awaited __anext__\(\)/asend\(\)"
1143            ):
1144                await nxt
1145
1146            await it.aclose()  # prevent unfinished iterator warning
1147
1148        self.loop.run_until_complete(run())
1149
1150    def test_async_gen_await_same_aclose_coro_twice(self):
1151        async def async_iterate():
1152            yield 1
1153            yield 2
1154
1155        async def run():
1156            it = async_iterate()
1157            nxt = it.aclose()
1158            await nxt
1159            with self.assertRaisesRegex(
1160                    RuntimeError,
1161                    r"cannot reuse already awaited aclose\(\)/athrow\(\)"
1162            ):
1163                await nxt
1164
1165        self.loop.run_until_complete(run())
1166
1167    def test_async_gen_aclose_twice_with_different_coros(self):
1168        # Regression test for https://bugs.python.org/issue39606
1169        async def async_iterate():
1170            yield 1
1171            yield 2
1172
1173        async def run():
1174            it = async_iterate()
1175            await it.aclose()
1176            await it.aclose()
1177
1178        self.loop.run_until_complete(run())
1179
1180    def test_async_gen_aclose_after_exhaustion(self):
1181        # Regression test for https://bugs.python.org/issue39606
1182        async def async_iterate():
1183            yield 1
1184            yield 2
1185
1186        async def run():
1187            it = async_iterate()
1188            async for _ in it:
1189                pass
1190            await it.aclose()
1191
1192        self.loop.run_until_complete(run())
1193
1194    def test_async_gen_aclose_compatible_with_get_stack(self):
1195        async def async_generator():
1196            yield object()
1197
1198        async def run():
1199            ag = async_generator()
1200            asyncio.create_task(ag.aclose())
1201            tasks = asyncio.all_tasks()
1202            for task in tasks:
1203                # No AttributeError raised
1204                task.get_stack()
1205
1206        self.loop.run_until_complete(run())
1207
1208
1209if __name__ == "__main__":
1210    unittest.main()
1211