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