1__all__ = ('Runner', 'run')
2
3import contextvars
4import enum
5import functools
6import threading
7import signal
8import sys
9from . import coroutines
10from . import events
11from . import exceptions
12from . import tasks
13
14
15class _State(enum.Enum):
16    CREATED = "created"
17    INITIALIZED = "initialized"
18    CLOSED = "closed"
19
20
21class Runner:
22    """A context manager that controls event loop life cycle.
23
24    The context manager always creates a new event loop,
25    allows to run async functions inside it,
26    and properly finalizes the loop at the context manager exit.
27
28    If debug is True, the event loop will be run in debug mode.
29    If loop_factory is passed, it is used for new event loop creation.
30
31    asyncio.run(main(), debug=True)
32
33    is a shortcut for
34
35    with asyncio.Runner(debug=True) as runner:
36        runner.run(main())
37
38    The run() method can be called multiple times within the runner's context.
39
40    This can be useful for interactive console (e.g. IPython),
41    unittest runners, console tools, -- everywhere when async code
42    is called from existing sync framework and where the preferred single
43    asyncio.run() call doesn't work.
44
45    """
46
47    # Note: the class is final, it is not intended for inheritance.
48
49    def __init__(self, *, debug=None, loop_factory=None):
50        self._state = _State.CREATED
51        self._debug = debug
52        self._loop_factory = loop_factory
53        self._loop = None
54        self._context = None
55        self._interrupt_count = 0
56        self._set_event_loop = False
57
58    def __enter__(self):
59        self._lazy_init()
60        return self
61
62    def __exit__(self, exc_type, exc_val, exc_tb):
63        self.close()
64
65    def close(self):
66        """Shutdown and close event loop."""
67        if self._state is not _State.INITIALIZED:
68            return
69        try:
70            loop = self._loop
71            _cancel_all_tasks(loop)
72            loop.run_until_complete(loop.shutdown_asyncgens())
73            loop.run_until_complete(loop.shutdown_default_executor())
74        finally:
75            if self._set_event_loop:
76                events.set_event_loop(None)
77            loop.close()
78            self._loop = None
79            self._state = _State.CLOSED
80
81    def get_loop(self):
82        """Return embedded event loop."""
83        self._lazy_init()
84        return self._loop
85
86    def run(self, coro, *, context=None):
87        """Run a coroutine inside the embedded event loop."""
88        if not coroutines.iscoroutine(coro):
89            raise ValueError("a coroutine was expected, got {!r}".format(coro))
90
91        if events._get_running_loop() is not None:
92            # fail fast with short traceback
93            raise RuntimeError(
94                "Runner.run() cannot be called from a running event loop")
95
96        self._lazy_init()
97
98        if context is None:
99            context = self._context
100        task = self._loop.create_task(coro, context=context)
101
102        if (threading.current_thread() is threading.main_thread()
103            and signal.getsignal(signal.SIGINT) is signal.default_int_handler
104        ):
105            sigint_handler = functools.partial(self._on_sigint, main_task=task)
106            try:
107                signal.signal(signal.SIGINT, sigint_handler)
108            except ValueError:
109                # `signal.signal` may throw if `threading.main_thread` does
110                # not support signals (e.g. embedded interpreter with signals
111                # not registered - see gh-91880)
112                sigint_handler = None
113        else:
114            sigint_handler = None
115
116        self._interrupt_count = 0
117        try:
118            return self._loop.run_until_complete(task)
119        except exceptions.CancelledError:
120            if self._interrupt_count > 0:
121                uncancel = getattr(task, "uncancel", None)
122                if uncancel is not None and uncancel() == 0:
123                    raise KeyboardInterrupt()
124            raise  # CancelledError
125        finally:
126            if (sigint_handler is not None
127                and signal.getsignal(signal.SIGINT) is sigint_handler
128            ):
129                signal.signal(signal.SIGINT, signal.default_int_handler)
130
131    def _lazy_init(self):
132        if self._state is _State.CLOSED:
133            raise RuntimeError("Runner is closed")
134        if self._state is _State.INITIALIZED:
135            return
136        if self._loop_factory is None:
137            self._loop = events.new_event_loop()
138            if not self._set_event_loop:
139                # Call set_event_loop only once to avoid calling
140                # attach_loop multiple times on child watchers
141                events.set_event_loop(self._loop)
142                self._set_event_loop = True
143        else:
144            self._loop = self._loop_factory()
145        if self._debug is not None:
146            self._loop.set_debug(self._debug)
147        self._context = contextvars.copy_context()
148        self._state = _State.INITIALIZED
149
150    def _on_sigint(self, signum, frame, main_task):
151        self._interrupt_count += 1
152        if self._interrupt_count == 1 and not main_task.done():
153            main_task.cancel()
154            # wakeup loop if it is blocked by select() with long timeout
155            self._loop.call_soon_threadsafe(lambda: None)
156            return
157        raise KeyboardInterrupt()
158
159
160def run(main, *, debug=None):
161    """Execute the coroutine and return the result.
162
163    This function runs the passed coroutine, taking care of
164    managing the asyncio event loop and finalizing asynchronous
165    generators.
166
167    This function cannot be called when another asyncio event loop is
168    running in the same thread.
169
170    If debug is True, the event loop will be run in debug mode.
171
172    This function always creates a new event loop and closes it at the end.
173    It should be used as a main entry point for asyncio programs, and should
174    ideally only be called once.
175
176    Example:
177
178        async def main():
179            await asyncio.sleep(1)
180            print('hello')
181
182        asyncio.run(main())
183    """
184    if events._get_running_loop() is not None:
185        # fail fast with short traceback
186        raise RuntimeError(
187            "asyncio.run() cannot be called from a running event loop")
188
189    with Runner(debug=debug) as runner:
190        return runner.run(main)
191
192
193def _cancel_all_tasks(loop):
194    to_cancel = tasks.all_tasks(loop)
195    if not to_cancel:
196        return
197
198    for task in to_cancel:
199        task.cancel()
200
201    loop.run_until_complete(tasks.gather(*to_cancel, return_exceptions=True))
202
203    for task in to_cancel:
204        if task.cancelled():
205            continue
206        if task.exception() is not None:
207            loop.call_exception_handler({
208                'message': 'unhandled exception during asyncio.run() shutdown',
209                'exception': task.exception(),
210                'task': task,
211            })
212