blob: 99464e3d0929fdbdb0c5eef5a116b647c4b50045 [file] [log] [blame]
Yury Selivanoveb636452016-09-08 22:01:51 -07001import inspect
Yury Selivanoveb636452016-09-08 22:01:51 -07002import types
3import unittest
4
Hai Shi46605972020-08-04 00:49:18 +08005from test.support.import_helper import import_module
Martin Panter94332cb2016-10-20 05:10:44 +00006asyncio = import_module("asyncio")
7
Yury Selivanoveb636452016-09-08 22:01:51 -07008
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:
Yury Selivanov52698c72018-06-07 20:31:26 -0400114 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
Yury Selivanoveb636452016-09-08 22:01:51 -0700136 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__()
Yury Selivanovfc4a0442019-09-29 22:59:11 -0700161
162 an = ai.__anext__()
163 self.assertEqual(an.__next__(), ('result',))
Yury Selivanoveb636452016-09-08 22:01:51 -0700164
165 try:
Yury Selivanovfc4a0442019-09-29 22:59:11 -0700166 an.__next__()
Yury Selivanoveb636452016-09-08 22:01:51 -0700167 except StopIteration as ex:
168 self.assertEqual(ex.args[0], 123)
169 else:
170 self.fail('StopIteration was not raised')
171
Yury Selivanovfc4a0442019-09-29 22:59:11 -0700172 an = ai.__anext__()
173 self.assertEqual(an.__next__(), ('result',))
Yury Selivanoveb636452016-09-08 22:01:51 -0700174
175 try:
Yury Selivanovfc4a0442019-09-29 22:59:11 -0700176 an.__next__()
Yury Selivanoveb636452016-09-08 22:01:51 -0700177 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__()
Yury Selivanovfc4a0442019-09-29 22:59:11 -0700200 an = ai.__anext__()
201 self.assertEqual(an.__next__(), ('result',))
Yury Selivanoveb636452016-09-08 22:01:51 -0700202
203 try:
Yury Selivanovfc4a0442019-09-29 22:59:11 -0700204 an.__next__()
Yury Selivanoveb636452016-09-08 22:01:51 -0700205 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
Yury Selivanov52698c72018-06-07 20:31:26 -0400311 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
Yury Selivanoveb636452016-09-08 22:01:51 -0700342 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
Brett Cannon8425de42018-06-01 20:34:09 -0700373 asyncio.set_event_loop_policy(None)
Yury Selivanoveb636452016-09-08 22:01:51 -0700374
Joshua Bronsonf0a6fde2021-03-23 18:47:21 -0400375 def test_async_gen_anext(self):
376 async def gen():
377 yield 1
378 yield 2
379 g = gen()
380 async def consume():
381 results = []
382 results.append(await anext(g))
383 results.append(await anext(g))
384 results.append(await anext(g, 'buckle my shoe'))
385 return results
386 res = self.loop.run_until_complete(consume())
387 self.assertEqual(res, [1, 2, 'buckle my shoe'])
388 with self.assertRaises(StopAsyncIteration):
389 self.loop.run_until_complete(consume())
390
391 def test_async_gen_aiter(self):
392 async def gen():
393 yield 1
394 yield 2
395 g = gen()
396 async def consume():
397 return [i async for i in aiter(g)]
398 res = self.loop.run_until_complete(consume())
399 self.assertEqual(res, [1, 2])
400
401 def test_async_gen_aiter_class(self):
402 results = []
403 class Gen:
404 async def __aiter__(self):
405 yield 1
406 yield 2
407 g = Gen()
408 async def consume():
409 ait = aiter(g)
410 while True:
411 try:
412 results.append(await anext(ait))
413 except StopAsyncIteration:
414 break
415 self.loop.run_until_complete(consume())
416 self.assertEqual(results, [1, 2])
417
418 def test_aiter_idempotent(self):
419 async def gen():
420 yield 1
421 applied_once = aiter(gen())
422 applied_twice = aiter(applied_once)
423 self.assertIs(applied_once, applied_twice)
424
425 def test_anext_bad_args(self):
426 async def gen():
427 yield 1
428 async def call_with_too_few_args():
429 await anext()
430 async def call_with_too_many_args():
431 await anext(gen(), 1, 3)
432 async def call_with_wrong_type_args():
433 await anext(1, gen())
434 with self.assertRaises(TypeError):
435 self.loop.run_until_complete(call_with_too_few_args())
436 with self.assertRaises(TypeError):
437 self.loop.run_until_complete(call_with_too_many_args())
438 with self.assertRaises(TypeError):
439 self.loop.run_until_complete(call_with_wrong_type_args())
440
441 def test_aiter_bad_args(self):
442 async def gen():
443 yield 1
444 async def call_with_too_few_args():
445 await aiter()
446 async def call_with_too_many_args():
447 await aiter(gen(), 1)
448 async def call_with_wrong_type_arg():
449 await aiter(1)
450 with self.assertRaises(TypeError):
451 self.loop.run_until_complete(call_with_too_few_args())
452 with self.assertRaises(TypeError):
453 self.loop.run_until_complete(call_with_too_many_args())
454 with self.assertRaises(TypeError):
455 self.loop.run_until_complete(call_with_wrong_type_arg())
456
Yury Selivanoveb636452016-09-08 22:01:51 -0700457 async def to_list(self, gen):
458 res = []
459 async for i in gen:
460 res.append(i)
461 return res
462
463 def test_async_gen_asyncio_01(self):
464 async def gen():
465 yield 1
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400466 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700467 yield 2
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400468 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700469 return
470 yield 3
471
472 res = self.loop.run_until_complete(self.to_list(gen()))
473 self.assertEqual(res, [1, 2])
474
475 def test_async_gen_asyncio_02(self):
476 async def gen():
477 yield 1
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400478 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700479 yield 2
480 1 / 0
481 yield 3
482
483 with self.assertRaises(ZeroDivisionError):
484 self.loop.run_until_complete(self.to_list(gen()))
485
486 def test_async_gen_asyncio_03(self):
487 loop = self.loop
488
489 class Gen:
490 async def __aiter__(self):
491 yield 1
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400492 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700493 yield 2
494
495 res = loop.run_until_complete(self.to_list(Gen()))
496 self.assertEqual(res, [1, 2])
497
498 def test_async_gen_asyncio_anext_04(self):
499 async def foo():
500 yield 1
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400501 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700502 try:
503 yield 2
504 yield 3
505 except ZeroDivisionError:
506 yield 1000
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400507 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700508 yield 4
509
510 async def run1():
511 it = foo().__aiter__()
512
513 self.assertEqual(await it.__anext__(), 1)
514 self.assertEqual(await it.__anext__(), 2)
515 self.assertEqual(await it.__anext__(), 3)
516 self.assertEqual(await it.__anext__(), 4)
517 with self.assertRaises(StopAsyncIteration):
518 await it.__anext__()
519 with self.assertRaises(StopAsyncIteration):
520 await it.__anext__()
521
522 async def run2():
523 it = foo().__aiter__()
524
525 self.assertEqual(await it.__anext__(), 1)
526 self.assertEqual(await it.__anext__(), 2)
527 try:
528 it.__anext__().throw(ZeroDivisionError)
529 except StopIteration as ex:
530 self.assertEqual(ex.args[0], 1000)
531 else:
532 self.fail('StopIteration was not raised')
533 self.assertEqual(await it.__anext__(), 4)
534 with self.assertRaises(StopAsyncIteration):
535 await it.__anext__()
536
537 self.loop.run_until_complete(run1())
538 self.loop.run_until_complete(run2())
539
540 def test_async_gen_asyncio_anext_05(self):
541 async def foo():
542 v = yield 1
543 v = yield v
544 yield v * 100
545
546 async def run():
547 it = foo().__aiter__()
548
549 try:
550 it.__anext__().send(None)
551 except StopIteration as ex:
552 self.assertEqual(ex.args[0], 1)
553 else:
554 self.fail('StopIteration was not raised')
555
556 try:
557 it.__anext__().send(10)
558 except StopIteration as ex:
559 self.assertEqual(ex.args[0], 10)
560 else:
561 self.fail('StopIteration was not raised')
562
563 try:
564 it.__anext__().send(12)
565 except StopIteration as ex:
566 self.assertEqual(ex.args[0], 1200)
567 else:
568 self.fail('StopIteration was not raised')
569
570 with self.assertRaises(StopAsyncIteration):
571 await it.__anext__()
572
573 self.loop.run_until_complete(run())
574
Yury Selivanov41782e42016-11-16 18:16:17 -0500575 def test_async_gen_asyncio_anext_06(self):
576 DONE = 0
577
578 # test synchronous generators
579 def foo():
580 try:
581 yield
582 except:
583 pass
584 g = foo()
585 g.send(None)
586 with self.assertRaises(StopIteration):
587 g.send(None)
588
589 # now with asynchronous generators
590
591 async def gen():
592 nonlocal DONE
593 try:
594 yield
595 except:
596 pass
597 DONE = 1
598
599 async def run():
600 nonlocal DONE
601 g = gen()
602 await g.asend(None)
603 with self.assertRaises(StopAsyncIteration):
604 await g.asend(None)
605 DONE += 10
606
607 self.loop.run_until_complete(run())
608 self.assertEqual(DONE, 11)
609
Serhiy Storchaka60e49aa2016-11-06 18:47:03 +0200610 def test_async_gen_asyncio_anext_tuple(self):
611 async def foo():
612 try:
613 yield (1,)
614 except ZeroDivisionError:
615 yield (2,)
616
617 async def run():
618 it = foo().__aiter__()
619
620 self.assertEqual(await it.__anext__(), (1,))
621 with self.assertRaises(StopIteration) as cm:
622 it.__anext__().throw(ZeroDivisionError)
623 self.assertEqual(cm.exception.args[0], (2,))
624 with self.assertRaises(StopAsyncIteration):
625 await it.__anext__()
626
627 self.loop.run_until_complete(run())
628
629 def test_async_gen_asyncio_anext_stopiteration(self):
630 async def foo():
631 try:
632 yield StopIteration(1)
633 except ZeroDivisionError:
634 yield StopIteration(3)
635
636 async def run():
637 it = foo().__aiter__()
638
639 v = await it.__anext__()
640 self.assertIsInstance(v, StopIteration)
641 self.assertEqual(v.value, 1)
642 with self.assertRaises(StopIteration) as cm:
643 it.__anext__().throw(ZeroDivisionError)
644 v = cm.exception.args[0]
645 self.assertIsInstance(v, StopIteration)
646 self.assertEqual(v.value, 3)
647 with self.assertRaises(StopAsyncIteration):
648 await it.__anext__()
649
650 self.loop.run_until_complete(run())
651
Yury Selivanoveb636452016-09-08 22:01:51 -0700652 def test_async_gen_asyncio_aclose_06(self):
653 async def foo():
654 try:
655 yield 1
656 1 / 0
657 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400658 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700659 yield 12
660
661 async def run():
662 gen = foo()
663 it = gen.__aiter__()
664 await it.__anext__()
665 await gen.aclose()
666
667 with self.assertRaisesRegex(
668 RuntimeError,
669 "async generator ignored GeneratorExit"):
670 self.loop.run_until_complete(run())
671
672 def test_async_gen_asyncio_aclose_07(self):
673 DONE = 0
674
675 async def foo():
676 nonlocal DONE
677 try:
678 yield 1
679 1 / 0
680 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400681 await asyncio.sleep(0.01)
682 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700683 DONE += 1
684 DONE += 1000
685
686 async def run():
687 gen = foo()
688 it = gen.__aiter__()
689 await it.__anext__()
690 await gen.aclose()
691
692 self.loop.run_until_complete(run())
693 self.assertEqual(DONE, 1)
694
695 def test_async_gen_asyncio_aclose_08(self):
696 DONE = 0
697
698 fut = asyncio.Future(loop=self.loop)
699
700 async def foo():
701 nonlocal DONE
702 try:
703 yield 1
704 await fut
705 DONE += 1000
706 yield 2
707 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400708 await asyncio.sleep(0.01)
709 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700710 DONE += 1
711 DONE += 1000
712
713 async def run():
714 gen = foo()
715 it = gen.__aiter__()
716 self.assertEqual(await it.__anext__(), 1)
Yury Selivanoveb636452016-09-08 22:01:51 -0700717 await gen.aclose()
Yury Selivanoveb636452016-09-08 22:01:51 -0700718
Yury Selivanovfc4a0442019-09-29 22:59:11 -0700719 self.loop.run_until_complete(run())
Yury Selivanoveb636452016-09-08 22:01:51 -0700720 self.assertEqual(DONE, 1)
721
722 # Silence ResourceWarnings
723 fut.cancel()
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400724 self.loop.run_until_complete(asyncio.sleep(0.01))
Yury Selivanoveb636452016-09-08 22:01:51 -0700725
726 def test_async_gen_asyncio_gc_aclose_09(self):
727 DONE = 0
728
729 async def gen():
730 nonlocal DONE
731 try:
732 while True:
733 yield 1
734 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400735 await asyncio.sleep(0.01)
736 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700737 DONE = 1
738
739 async def run():
740 g = gen()
741 await g.__anext__()
742 await g.__anext__()
743 del g
744
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400745 await asyncio.sleep(0.1)
Yury Selivanoveb636452016-09-08 22:01:51 -0700746
747 self.loop.run_until_complete(run())
748 self.assertEqual(DONE, 1)
749
Yury Selivanov41782e42016-11-16 18:16:17 -0500750 def test_async_gen_asyncio_aclose_10(self):
751 DONE = 0
752
753 # test synchronous generators
754 def foo():
755 try:
756 yield
757 except:
758 pass
759 g = foo()
760 g.send(None)
761 g.close()
762
763 # now with asynchronous generators
764
765 async def gen():
766 nonlocal DONE
767 try:
768 yield
769 except:
770 pass
771 DONE = 1
772
773 async def run():
774 nonlocal DONE
775 g = gen()
776 await g.asend(None)
777 await g.aclose()
778 DONE += 10
779
780 self.loop.run_until_complete(run())
781 self.assertEqual(DONE, 11)
782
783 def test_async_gen_asyncio_aclose_11(self):
784 DONE = 0
785
786 # test synchronous generators
787 def foo():
788 try:
789 yield
790 except:
791 pass
792 yield
793 g = foo()
794 g.send(None)
795 with self.assertRaisesRegex(RuntimeError, 'ignored GeneratorExit'):
796 g.close()
797
798 # now with asynchronous generators
799
800 async def gen():
801 nonlocal DONE
802 try:
803 yield
804 except:
805 pass
806 yield
807 DONE += 1
808
809 async def run():
810 nonlocal DONE
811 g = gen()
812 await g.asend(None)
813 with self.assertRaisesRegex(RuntimeError, 'ignored GeneratorExit'):
814 await g.aclose()
815 DONE += 10
816
817 self.loop.run_until_complete(run())
818 self.assertEqual(DONE, 10)
819
Vincent Michel8e0de2a2019-11-19 05:53:52 -0800820 def test_async_gen_asyncio_aclose_12(self):
821 DONE = 0
822
823 async def target():
824 await asyncio.sleep(0.01)
825 1 / 0
826
827 async def foo():
828 nonlocal DONE
829 task = asyncio.create_task(target())
830 try:
831 yield 1
832 finally:
833 try:
834 await task
835 except ZeroDivisionError:
836 DONE = 1
837
838 async def run():
839 gen = foo()
840 it = gen.__aiter__()
841 await it.__anext__()
842 await gen.aclose()
843
844 self.loop.run_until_complete(run())
845 self.assertEqual(DONE, 1)
846
Yury Selivanoveb636452016-09-08 22:01:51 -0700847 def test_async_gen_asyncio_asend_01(self):
848 DONE = 0
849
850 # Sanity check:
851 def sgen():
852 v = yield 1
853 yield v * 2
854 sg = sgen()
855 v = sg.send(None)
856 self.assertEqual(v, 1)
857 v = sg.send(100)
858 self.assertEqual(v, 200)
859
860 async def gen():
861 nonlocal DONE
862 try:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400863 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700864 v = yield 1
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400865 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700866 yield v * 2
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400867 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700868 return
869 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400870 await asyncio.sleep(0.01)
871 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700872 DONE = 1
873
874 async def run():
875 g = gen()
876
877 v = await g.asend(None)
878 self.assertEqual(v, 1)
879
880 v = await g.asend(100)
881 self.assertEqual(v, 200)
882
883 with self.assertRaises(StopAsyncIteration):
884 await g.asend(None)
885
886 self.loop.run_until_complete(run())
887 self.assertEqual(DONE, 1)
888
889 def test_async_gen_asyncio_asend_02(self):
890 DONE = 0
891
892 async def sleep_n_crash(delay):
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400893 await asyncio.sleep(delay)
Yury Selivanoveb636452016-09-08 22:01:51 -0700894 1 / 0
895
896 async def gen():
897 nonlocal DONE
898 try:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400899 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700900 v = yield 1
901 await sleep_n_crash(0.01)
902 DONE += 1000
903 yield v * 2
904 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400905 await asyncio.sleep(0.01)
906 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700907 DONE = 1
908
909 async def run():
910 g = gen()
911
912 v = await g.asend(None)
913 self.assertEqual(v, 1)
914
915 await g.asend(100)
916
917 with self.assertRaises(ZeroDivisionError):
918 self.loop.run_until_complete(run())
919 self.assertEqual(DONE, 1)
920
921 def test_async_gen_asyncio_asend_03(self):
922 DONE = 0
923
924 async def sleep_n_crash(delay):
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400925 fut = asyncio.ensure_future(asyncio.sleep(delay),
Yury Selivanoveb636452016-09-08 22:01:51 -0700926 loop=self.loop)
927 self.loop.call_later(delay / 2, lambda: fut.cancel())
928 return await fut
929
930 async def gen():
931 nonlocal DONE
932 try:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400933 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700934 v = yield 1
935 await sleep_n_crash(0.01)
936 DONE += 1000
937 yield v * 2
938 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400939 await asyncio.sleep(0.01)
940 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700941 DONE = 1
942
943 async def run():
944 g = gen()
945
946 v = await g.asend(None)
947 self.assertEqual(v, 1)
948
949 await g.asend(100)
950
951 with self.assertRaises(asyncio.CancelledError):
952 self.loop.run_until_complete(run())
953 self.assertEqual(DONE, 1)
954
955 def test_async_gen_asyncio_athrow_01(self):
956 DONE = 0
957
958 class FooEr(Exception):
959 pass
960
961 # Sanity check:
962 def sgen():
963 try:
964 v = yield 1
965 except FooEr:
966 v = 1000
967 yield v * 2
968 sg = sgen()
969 v = sg.send(None)
970 self.assertEqual(v, 1)
971 v = sg.throw(FooEr)
972 self.assertEqual(v, 2000)
973 with self.assertRaises(StopIteration):
974 sg.send(None)
975
976 async def gen():
977 nonlocal DONE
978 try:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400979 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700980 try:
981 v = yield 1
982 except FooEr:
983 v = 1000
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400984 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700985 yield v * 2
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400986 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700987 # return
988 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -0400989 await asyncio.sleep(0.01)
990 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -0700991 DONE = 1
992
993 async def run():
994 g = gen()
995
996 v = await g.asend(None)
997 self.assertEqual(v, 1)
998
999 v = await g.athrow(FooEr)
1000 self.assertEqual(v, 2000)
1001
1002 with self.assertRaises(StopAsyncIteration):
1003 await g.asend(None)
1004
1005 self.loop.run_until_complete(run())
1006 self.assertEqual(DONE, 1)
1007
1008 def test_async_gen_asyncio_athrow_02(self):
1009 DONE = 0
1010
1011 class FooEr(Exception):
1012 pass
1013
1014 async def sleep_n_crash(delay):
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001015 fut = asyncio.ensure_future(asyncio.sleep(delay),
Yury Selivanoveb636452016-09-08 22:01:51 -07001016 loop=self.loop)
1017 self.loop.call_later(delay / 2, lambda: fut.cancel())
1018 return await fut
1019
1020 async def gen():
1021 nonlocal DONE
1022 try:
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001023 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -07001024 try:
1025 v = yield 1
1026 except FooEr:
1027 await sleep_n_crash(0.01)
1028 yield v * 2
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001029 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -07001030 # return
1031 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001032 await asyncio.sleep(0.01)
1033 await asyncio.sleep(0.01)
Yury Selivanoveb636452016-09-08 22:01:51 -07001034 DONE = 1
1035
1036 async def run():
1037 g = gen()
1038
1039 v = await g.asend(None)
1040 self.assertEqual(v, 1)
1041
1042 try:
1043 await g.athrow(FooEr)
1044 except asyncio.CancelledError:
1045 self.assertEqual(DONE, 1)
1046 raise
1047 else:
1048 self.fail('CancelledError was not raised')
1049
1050 with self.assertRaises(asyncio.CancelledError):
1051 self.loop.run_until_complete(run())
1052 self.assertEqual(DONE, 1)
1053
Yury Selivanov41782e42016-11-16 18:16:17 -05001054 def test_async_gen_asyncio_athrow_03(self):
1055 DONE = 0
1056
1057 # test synchronous generators
1058 def foo():
1059 try:
1060 yield
1061 except:
1062 pass
1063 g = foo()
1064 g.send(None)
1065 with self.assertRaises(StopIteration):
1066 g.throw(ValueError)
1067
1068 # now with asynchronous generators
1069
1070 async def gen():
1071 nonlocal DONE
1072 try:
1073 yield
1074 except:
1075 pass
1076 DONE = 1
1077
1078 async def run():
1079 nonlocal DONE
1080 g = gen()
1081 await g.asend(None)
1082 with self.assertRaises(StopAsyncIteration):
1083 await g.athrow(ValueError)
1084 DONE += 10
1085
1086 self.loop.run_until_complete(run())
1087 self.assertEqual(DONE, 11)
1088
Serhiy Storchaka60e49aa2016-11-06 18:47:03 +02001089 def test_async_gen_asyncio_athrow_tuple(self):
1090 async def gen():
1091 try:
1092 yield 1
1093 except ZeroDivisionError:
1094 yield (2,)
1095
1096 async def run():
1097 g = gen()
1098 v = await g.asend(None)
1099 self.assertEqual(v, 1)
1100 v = await g.athrow(ZeroDivisionError)
1101 self.assertEqual(v, (2,))
1102 with self.assertRaises(StopAsyncIteration):
1103 await g.asend(None)
1104
1105 self.loop.run_until_complete(run())
1106
1107 def test_async_gen_asyncio_athrow_stopiteration(self):
1108 async def gen():
1109 try:
1110 yield 1
1111 except ZeroDivisionError:
1112 yield StopIteration(2)
1113
1114 async def run():
1115 g = gen()
1116 v = await g.asend(None)
1117 self.assertEqual(v, 1)
1118 v = await g.athrow(ZeroDivisionError)
1119 self.assertIsInstance(v, StopIteration)
1120 self.assertEqual(v.value, 2)
1121 with self.assertRaises(StopAsyncIteration):
1122 await g.asend(None)
1123
1124 self.loop.run_until_complete(run())
1125
Yury Selivanoveb636452016-09-08 22:01:51 -07001126 def test_async_gen_asyncio_shutdown_01(self):
1127 finalized = 0
1128
1129 async def waiter(timeout):
1130 nonlocal finalized
1131 try:
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001132 await asyncio.sleep(timeout)
Yury Selivanoveb636452016-09-08 22:01:51 -07001133 yield 1
1134 finally:
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001135 await asyncio.sleep(0)
Yury Selivanoveb636452016-09-08 22:01:51 -07001136 finalized += 1
1137
1138 async def wait():
1139 async for _ in waiter(1):
1140 pass
1141
1142 t1 = self.loop.create_task(wait())
1143 t2 = self.loop.create_task(wait())
1144
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001145 self.loop.run_until_complete(asyncio.sleep(0.1))
Yury Selivanoveb636452016-09-08 22:01:51 -07001146
Yury Selivanoveb636452016-09-08 22:01:51 -07001147 # Silence warnings
1148 t1.cancel()
1149 t2.cancel()
Yury Selivanoveb636452016-09-08 22:01:51 -07001150
Yury Selivanovfc4a0442019-09-29 22:59:11 -07001151 with self.assertRaises(asyncio.CancelledError):
1152 self.loop.run_until_complete(t1)
1153 with self.assertRaises(asyncio.CancelledError):
1154 self.loop.run_until_complete(t2)
Yury Selivanoveb636452016-09-08 22:01:51 -07001155
Yury Selivanoveb636452016-09-08 22:01:51 -07001156 self.loop.run_until_complete(self.loop.shutdown_asyncgens())
1157
Yury Selivanovfc4a0442019-09-29 22:59:11 -07001158 self.assertEqual(finalized, 2)
Yury Selivanoveb636452016-09-08 22:01:51 -07001159
Yury Selivanovb8ab9d32017-10-06 02:58:28 -04001160 def test_async_gen_expression_01(self):
1161 async def arange(n):
1162 for i in range(n):
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001163 await asyncio.sleep(0.01)
Yury Selivanovb8ab9d32017-10-06 02:58:28 -04001164 yield i
1165
1166 def make_arange(n):
1167 # This syntax is legal starting with Python 3.7
1168 return (i * 2 async for i in arange(n))
1169
1170 async def run():
1171 return [i async for i in make_arange(10)]
1172
1173 res = self.loop.run_until_complete(run())
1174 self.assertEqual(res, [i * 2 for i in range(10)])
1175
1176 def test_async_gen_expression_02(self):
1177 async def wrap(n):
Yury Selivanov9012a0f2018-10-02 13:53:06 -04001178 await asyncio.sleep(0.01)
Yury Selivanovb8ab9d32017-10-06 02:58:28 -04001179 return n
1180
1181 def make_arange(n):
1182 # This syntax is legal starting with Python 3.7
1183 return (i * 2 for i in range(n) if await wrap(i))
1184
1185 async def run():
1186 return [i async for i in make_arange(10)]
1187
1188 res = self.loop.run_until_complete(run())
1189 self.assertEqual(res, [i * 2 for i in range(1, 10)])
1190
Andrew Svetlovc2753122019-09-17 15:59:49 +03001191 def test_asyncgen_nonstarted_hooks_are_cancellable(self):
1192 # See https://bugs.python.org/issue38013
1193 messages = []
1194
1195 def exception_handler(loop, context):
1196 messages.append(context)
1197
1198 async def async_iterate():
1199 yield 1
1200 yield 2
1201
1202 async def main():
1203 loop = asyncio.get_running_loop()
1204 loop.set_exception_handler(exception_handler)
1205
1206 async for i in async_iterate():
1207 break
1208
1209 asyncio.run(main())
1210
1211 self.assertEqual([], messages)
1212
Nathaniel J. Smith925dc7f2020-02-13 00:15:38 -08001213 def test_async_gen_await_same_anext_coro_twice(self):
Andrew Svetlova96e06d2020-01-21 00:49:30 +02001214 async def async_iterate():
1215 yield 1
1216 yield 2
1217
1218 async def run():
1219 it = async_iterate()
1220 nxt = it.__anext__()
1221 await nxt
1222 with self.assertRaisesRegex(
1223 RuntimeError,
1224 r"cannot reuse already awaited __anext__\(\)/asend\(\)"
1225 ):
1226 await nxt
1227
1228 await it.aclose() # prevent unfinished iterator warning
1229
1230 self.loop.run_until_complete(run())
1231
Nathaniel J. Smith925dc7f2020-02-13 00:15:38 -08001232 def test_async_gen_await_same_aclose_coro_twice(self):
Andrew Svetlova96e06d2020-01-21 00:49:30 +02001233 async def async_iterate():
1234 yield 1
1235 yield 2
1236
1237 async def run():
1238 it = async_iterate()
1239 nxt = it.aclose()
1240 await nxt
1241 with self.assertRaisesRegex(
1242 RuntimeError,
1243 r"cannot reuse already awaited aclose\(\)/athrow\(\)"
1244 ):
1245 await nxt
1246
1247 self.loop.run_until_complete(run())
1248
Nathaniel J. Smith925dc7f2020-02-13 00:15:38 -08001249 def test_async_gen_aclose_twice_with_different_coros(self):
1250 # Regression test for https://bugs.python.org/issue39606
1251 async def async_iterate():
1252 yield 1
1253 yield 2
1254
1255 async def run():
1256 it = async_iterate()
1257 await it.aclose()
1258 await it.aclose()
1259
1260 self.loop.run_until_complete(run())
1261
1262 def test_async_gen_aclose_after_exhaustion(self):
1263 # Regression test for https://bugs.python.org/issue39606
1264 async def async_iterate():
1265 yield 1
1266 yield 2
1267
1268 async def run():
1269 it = async_iterate()
1270 async for _ in it:
1271 pass
1272 await it.aclose()
1273
1274 self.loop.run_until_complete(run())
Yury Selivanovb8ab9d32017-10-06 02:58:28 -04001275
Lidi Zheng44823372020-03-02 04:45:54 -08001276 def test_async_gen_aclose_compatible_with_get_stack(self):
1277 async def async_generator():
1278 yield object()
1279
1280 async def run():
1281 ag = async_generator()
1282 asyncio.create_task(ag.aclose())
1283 tasks = asyncio.all_tasks()
1284 for task in tasks:
1285 # No AttributeError raised
1286 task.get_stack()
1287
1288 self.loop.run_until_complete(run())
1289
1290
Yury Selivanoveb636452016-09-08 22:01:51 -07001291if __name__ == "__main__":
1292 unittest.main()