|
| 1 | +importcontextvars |
| 2 | +importenum |
| 3 | +importfunctools |
1 | 4 | importos |
2 | 5 | importsignal |
3 | 6 | importsys |
| 7 | +importthreading |
4 | 8 | importunittest |
5 | 9 | importwarnings |
6 | 10 | fromunittestimportmock |
7 | 11 |
|
8 | | -importasyncio |
| 12 | +importasyncio.runners |
9 | 13 | fromasyncioimportbase_subprocess |
| 14 | +fromasyncioimportcoroutines |
| 15 | +fromasyncioimportevents |
10 | 16 | fromasyncioimportsubprocess |
11 | 17 | fromtest.test_asyncioimportutilsastest_utils |
12 | 18 | fromtestimportsupport |
@@ -636,6 +642,147 @@ async def execute(): |
636 | 642 | self.assertIsNone(self.loop.run_until_complete(execute())) |
637 | 643 |
|
638 | 644 |
|
| 645 | +class_State(enum.Enum): |
| 646 | +CREATED="created" |
| 647 | +INITIALIZED="initialized" |
| 648 | +CLOSED="closed" |
| 649 | + |
| 650 | +classRunner: |
| 651 | +"""A context manager that controls event loop life cycle. |
| 652 | +
|
| 653 | + The context manager always creates a new event loop, |
| 654 | + allows to run async functions inside it, |
| 655 | + and properly finalizes the loop at the context manager exit. |
| 656 | +
|
| 657 | + If debug is True, the event loop will be run in debug mode. |
| 658 | + If loop_factory is passed, it is used for new event loop creation. |
| 659 | +
|
| 660 | + asyncio.run(main(), debug=True) |
| 661 | +
|
| 662 | + is a shortcut for |
| 663 | +
|
| 664 | + with asyncio.Runner(debug=True) as runner: |
| 665 | + runner.run(main()) |
| 666 | +
|
| 667 | + The run() method can be called multiple times within the runner's context. |
| 668 | +
|
| 669 | + This can be useful for interactive console (e.g. IPython), |
| 670 | + unittest runners, console tools, -- everywhere when async code |
| 671 | + is called from existing sync framework and where the preferred single |
| 672 | + asyncio.run() call doesn't work. |
| 673 | +
|
| 674 | + """ |
| 675 | + |
| 676 | +# Note: the class is final, it is not intended for inheritance. |
| 677 | + |
| 678 | +def__init__(self,*,debug=None,loop_factory=None): |
| 679 | +self._state=_State.CREATED |
| 680 | +self._debug=debug |
| 681 | +self._loop_factory=loop_factory |
| 682 | +self._loop=None |
| 683 | +self._context=None |
| 684 | +self._interrupt_count=0 |
| 685 | +self._set_event_loop=False |
| 686 | + |
| 687 | +def__enter__(self): |
| 688 | +self._lazy_init() |
| 689 | +returnself |
| 690 | + |
| 691 | +def__exit__(self,exc_type,exc_val,exc_tb): |
| 692 | +self.close() |
| 693 | + |
| 694 | +defclose(self): |
| 695 | +"""Shutdown and close event loop.""" |
| 696 | +ifself._stateisnot_State.INITIALIZED: |
| 697 | +return |
| 698 | +try: |
| 699 | +loop=self._loop |
| 700 | +asyncio.runners._cancel_all_tasks(loop) |
| 701 | +loop.run_until_complete(loop.shutdown_asyncgens()) |
| 702 | +loop.run_until_complete(loop.shutdown_default_executor()) |
| 703 | +finally: |
| 704 | +ifself._set_event_loop: |
| 705 | +events.set_event_loop(None) |
| 706 | +loop.close() |
| 707 | +self._loop=None |
| 708 | +self._state=_State.CLOSED |
| 709 | + |
| 710 | +defget_loop(self): |
| 711 | +"""Return embedded event loop.""" |
| 712 | +self._lazy_init() |
| 713 | +returnself._loop |
| 714 | + |
| 715 | +defrun(self,coro,*,context=None): |
| 716 | +"""Run a coroutine inside the embedded event loop.""" |
| 717 | +ifnotcoroutines.iscoroutine(coro): |
| 718 | +raiseValueError("a coroutine was expected, got {!r}".format(coro)) |
| 719 | + |
| 720 | +ifevents._get_running_loop()isnotNone: |
| 721 | +# fail fast with short traceback |
| 722 | +raiseRuntimeError( |
| 723 | +"Runner.run() cannot be called from a running event loop") |
| 724 | + |
| 725 | +self._lazy_init() |
| 726 | + |
| 727 | +ifcontextisNone: |
| 728 | +context=self._context |
| 729 | +task=self._loop.create_task(coro) |
| 730 | + |
| 731 | +if (threading.current_thread()isthreading.main_thread() |
| 732 | +andsignal.getsignal(signal.SIGINT)issignal.default_int_handler |
| 733 | + ): |
| 734 | +sigint_handler=functools.partial(self._on_sigint,main_task=task) |
| 735 | +try: |
| 736 | +signal.signal(signal.SIGINT,sigint_handler) |
| 737 | +exceptValueError: |
| 738 | +# `signal.signal` may throw if `threading.main_thread` does |
| 739 | +# not support signals (e.g. embedded interpreter with signals |
| 740 | +# not registered - see gh-91880) |
| 741 | +sigint_handler=None |
| 742 | +else: |
| 743 | +sigint_handler=None |
| 744 | + |
| 745 | +self._interrupt_count=0 |
| 746 | +try: |
| 747 | +ifself._set_event_loop: |
| 748 | +events.set_event_loop(self._loop) |
| 749 | +returnself._loop.run_until_complete(task) |
| 750 | +exceptexceptions.CancelledError: |
| 751 | +ifself._interrupt_count>0andtask.uncancel()==0: |
| 752 | +raiseKeyboardInterrupt() |
| 753 | +else: |
| 754 | +raise# CancelledError |
| 755 | +finally: |
| 756 | +if (sigint_handlerisnotNone |
| 757 | +andsignal.getsignal(signal.SIGINT)issigint_handler |
| 758 | + ): |
| 759 | +signal.signal(signal.SIGINT,signal.default_int_handler) |
| 760 | + |
| 761 | +def_lazy_init(self): |
| 762 | +ifself._stateis_State.CLOSED: |
| 763 | +raiseRuntimeError("Runner is closed") |
| 764 | +ifself._stateis_State.INITIALIZED: |
| 765 | +return |
| 766 | +ifself._loop_factoryisNone: |
| 767 | +self._loop=events.new_event_loop() |
| 768 | +self._set_event_loop=True |
| 769 | +else: |
| 770 | +self._loop=self._loop_factory() |
| 771 | +ifself._debugisnotNone: |
| 772 | +self._loop.set_debug(self._debug) |
| 773 | +self._context=contextvars.copy_context() |
| 774 | +self._state=_State.INITIALIZED |
| 775 | + |
| 776 | +def_on_sigint(self,signum,frame,main_task): |
| 777 | +self._interrupt_count+=1 |
| 778 | +ifself._interrupt_count==1andnotmain_task.done(): |
| 779 | +main_task.cancel() |
| 780 | +# wakeup loop if it is blocked by select() with long timeout |
| 781 | +self._loop.call_soon_threadsafe(lambda:None) |
| 782 | +return |
| 783 | +raiseKeyboardInterrupt() |
| 784 | + |
| 785 | + |
639 | 786 | ifsys.platform!='win32': |
640 | 787 | # Unix |
641 | 788 | classSubprocessWatcherMixin(SubprocessMixin): |
@@ -717,7 +864,7 @@ async def execute(): |
717 | 864 |
|
718 | 865 | watcher.add_child_handler.assert_not_called() |
719 | 866 |
|
720 | | -withasyncio.Runner(loop_factory=asyncio.new_event_loop)asrunner: |
| 867 | +withRunner(loop_factory=asyncio.new_event_loop)asrunner: |
721 | 868 | self.assertIsNone(runner.run(execute())) |
722 | 869 | self.assertListEqual(watcher.mock_calls, [ |
723 | 870 | mock.call.__enter__(), |
|