blob: 717837d856843679c5b5601dccdcb693a250fb0b [file] [log] [blame]
Guido van Rossum27b7c7e2013-10-17 13:40:50 -07001"""Support for tasks, coroutines and the scheduler."""
2
Yury Selivanov6370f342017-12-10 18:36:12 -05003__all__ = (
Andrew Svetlovf74ef452017-12-15 07:04:38 +02004 'Task', 'create_task',
Yury Selivanov6370f342017-12-10 18:36:12 -05005 'FIRST_COMPLETED', 'FIRST_EXCEPTION', 'ALL_COMPLETED',
Yury Selivanov9edad3c2017-12-11 10:03:48 -05006 'wait', 'wait_for', 'as_completed', 'sleep',
Yury Selivanov6370f342017-12-10 18:36:12 -05007 'gather', 'shield', 'ensure_future', 'run_coroutine_threadsafe',
Andrew Svetlov44d1a592017-12-16 21:58:38 +02008 'current_task', 'all_tasks',
9 '_register_task', '_unregister_task', '_enter_task', '_leave_task',
Yury Selivanov6370f342017-12-10 18:36:12 -050010)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -070011
Guido van Rossum27b7c7e2013-10-17 13:40:50 -070012import concurrent.futures
Yury Selivanovf23746a2018-01-22 19:11:18 -050013import contextvars
Guido van Rossum27b7c7e2013-10-17 13:40:50 -070014import functools
15import inspect
Alex Grönholmcca4eec2018-08-09 00:06:47 +030016import itertools
Andrew Svetlov5f841b52017-12-09 00:23:48 +020017import types
Yury Selivanov59eb9a42015-05-11 14:48:38 -040018import warnings
Guido van Rossum27b7c7e2013-10-17 13:40:50 -070019import weakref
20
Yury Selivanova0c1ba62016-10-28 12:52:37 -040021from . import base_tasks
Victor Stinnerf951d282014-06-29 00:46:45 +020022from . import coroutines
Guido van Rossum27b7c7e2013-10-17 13:40:50 -070023from . import events
Andrew Svetlov0baa72f2018-09-11 10:13:04 -070024from . import exceptions
Guido van Rossum27b7c7e2013-10-17 13:40:50 -070025from . import futures
Andrew Svetlov68b34a72019-05-16 17:52:10 +030026from .coroutines import _is_coroutine
Guido van Rossum27b7c7e2013-10-17 13:40:50 -070027
Alex Grönholmcca4eec2018-08-09 00:06:47 +030028# Helper to generate new task names
29# This uses itertools.count() instead of a "+= 1" operation because the latter
30# is not thread safe. See bpo-11866 for a longer explanation.
31_task_name_counter = itertools.count(1).__next__
32
Guido van Rossum27b7c7e2013-10-17 13:40:50 -070033
Andrew Svetlov44d1a592017-12-16 21:58:38 +020034def current_task(loop=None):
35 """Return a currently executed task."""
36 if loop is None:
37 loop = events.get_running_loop()
38 return _current_tasks.get(loop)
39
40
41def all_tasks(loop=None):
42 """Return a set of all tasks for the loop."""
43 if loop is None:
Yury Selivanov416c1eb2018-05-28 17:54:02 -040044 loop = events.get_running_loop()
Andrew Svetlov65aa64f2019-06-11 18:27:30 +030045 # Looping over a WeakSet (_all_tasks) isn't safe as it can be updated from another
46 # thread while we do so. Therefore we cast it to list prior to filtering. The list
47 # cast itself requires iteration, so we repeat it several times ignoring
48 # RuntimeErrors (which are not very likely to occur). See issues 34970 and 36607 for
49 # details.
50 i = 0
51 while True:
52 try:
53 tasks = list(_all_tasks)
54 except RuntimeError:
55 i += 1
56 if i >= 1000:
57 raise
58 else:
59 break
60 return {t for t in tasks
Yury Selivanov416c1eb2018-05-28 17:54:02 -040061 if futures._get_loop(t) is loop and not t.done()}
62
63
64def _all_tasks_compat(loop=None):
65 # Different from "all_task()" by returning *all* Tasks, including
66 # the completed ones. Used to implement deprecated "Tasks.all_task()"
67 # method.
68 if loop is None:
Andrew Svetlov44d1a592017-12-16 21:58:38 +020069 loop = events.get_event_loop()
Andrew Svetlov65aa64f2019-06-11 18:27:30 +030070 # Looping over a WeakSet (_all_tasks) isn't safe as it can be updated from another
71 # thread while we do so. Therefore we cast it to list prior to filtering. The list
72 # cast itself requires iteration, so we repeat it several times ignoring
73 # RuntimeErrors (which are not very likely to occur). See issues 34970 and 36607 for
74 # details.
75 i = 0
76 while True:
77 try:
78 tasks = list(_all_tasks)
79 except RuntimeError:
80 i += 1
81 if i >= 1000:
82 raise
83 else:
84 break
85 return {t for t in tasks if futures._get_loop(t) is loop}
Andrew Svetlov44d1a592017-12-16 21:58:38 +020086
87
Alex Grönholmcca4eec2018-08-09 00:06:47 +030088def _set_task_name(task, name):
89 if name is not None:
90 try:
91 set_name = task.set_name
92 except AttributeError:
93 pass
94 else:
95 set_name(name)
96
97
Yury Selivanov0cf16f92017-12-25 10:48:15 -050098class Task(futures._PyFuture): # Inherit Python Task implementation
99 # from a Python Future implementation.
100
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700101 """A coroutine wrapped in a Future."""
102
103 # An important invariant maintained while a Task not done:
104 #
105 # - Either _fut_waiter is None, and _step() is scheduled;
106 # - or _fut_waiter is some Future, and _step() is *not* scheduled.
107 #
108 # The only transition from the latter to the former is through
109 # _wakeup(). When _fut_waiter is not None, one of its callbacks
110 # must be _wakeup().
111
Victor Stinnerfe22e092014-12-04 23:00:13 +0100112 # If False, don't log a message if the task is destroyed whereas its
113 # status is still pending
114 _log_destroy_pending = True
115
Guido van Rossum1a605ed2013-12-06 12:57:40 -0800116 @classmethod
117 def current_task(cls, loop=None):
118 """Return the currently running task in an event loop or None.
119
120 By default the current task for the current event loop is returned.
121
122 None is returned when called not in the context of a Task.
123 """
Matthias Bussonnierd0ebf132019-05-20 23:20:10 -0700124 warnings.warn("Task.current_task() is deprecated since Python 3.7, "
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200125 "use asyncio.current_task() instead",
Inada Naokic5c6cda2019-03-22 20:07:32 +0900126 DeprecationWarning,
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200127 stacklevel=2)
Guido van Rossum1a605ed2013-12-06 12:57:40 -0800128 if loop is None:
129 loop = events.get_event_loop()
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200130 return current_task(loop)
Guido van Rossum1a605ed2013-12-06 12:57:40 -0800131
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700132 @classmethod
133 def all_tasks(cls, loop=None):
134 """Return a set of all tasks for an event loop.
135
136 By default all tasks for the current event loop are returned.
137 """
Matthias Bussonnierd0ebf132019-05-20 23:20:10 -0700138 warnings.warn("Task.all_tasks() is deprecated since Python 3.7, "
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200139 "use asyncio.all_tasks() instead",
Inada Naokic5c6cda2019-03-22 20:07:32 +0900140 DeprecationWarning,
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200141 stacklevel=2)
Yury Selivanov416c1eb2018-05-28 17:54:02 -0400142 return _all_tasks_compat(loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700143
Alex Grönholmcca4eec2018-08-09 00:06:47 +0300144 def __init__(self, coro, *, loop=None, name=None):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700145 super().__init__(loop=loop)
Victor Stinner80f53aa2014-06-27 13:52:20 +0200146 if self._source_traceback:
147 del self._source_traceback[-1]
Andrew Svetlovf74ef452017-12-15 07:04:38 +0200148 if not coroutines.iscoroutine(coro):
149 # raise after Future.__init__(), attrs are required for __del__
150 # prevent logging for pending task in __del__
151 self._log_destroy_pending = False
152 raise TypeError(f"a coroutine was expected, got {coro!r}")
153
Alex Grönholmcca4eec2018-08-09 00:06:47 +0300154 if name is None:
155 self._name = f'Task-{_task_name_counter()}'
156 else:
157 self._name = str(name)
158
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700159 self._must_cancel = False
Andrew Svetlovf74ef452017-12-15 07:04:38 +0200160 self._fut_waiter = None
161 self._coro = coro
Yury Selivanovf23746a2018-01-22 19:11:18 -0500162 self._context = contextvars.copy_context()
Andrew Svetlovf74ef452017-12-15 07:04:38 +0200163
Yury Selivanov22feeb82018-01-24 11:31:01 -0500164 self._loop.call_soon(self.__step, context=self._context)
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500165 _register_task(self)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700166
INADA Naoki3e2ad8e2017-04-25 10:57:18 +0900167 def __del__(self):
168 if self._state == futures._PENDING and self._log_destroy_pending:
169 context = {
170 'task': self,
171 'message': 'Task was destroyed but it is pending!',
172 }
173 if self._source_traceback:
174 context['source_traceback'] = self._source_traceback
175 self._loop.call_exception_handler(context)
Yury Selivanov0cf16f92017-12-25 10:48:15 -0500176 super().__del__()
Victor Stinnera02f81f2014-06-24 22:37:53 +0200177
Batuhan Taşkayadec36722019-12-07 14:05:07 +0300178 def __class_getitem__(cls, type):
179 return cls
180
Victor Stinner313a9802014-07-29 12:58:23 +0200181 def _repr_info(self):
Yury Selivanova0c1ba62016-10-28 12:52:37 -0400182 return base_tasks._task_repr_info(self)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700183
Alex Grönholm98ef9202019-05-30 18:30:09 +0300184 def get_coro(self):
185 return self._coro
186
Alex Grönholmcca4eec2018-08-09 00:06:47 +0300187 def get_name(self):
188 return self._name
189
190 def set_name(self, value):
191 self._name = str(value)
192
Yury Selivanov0cf16f92017-12-25 10:48:15 -0500193 def set_result(self, result):
194 raise RuntimeError('Task does not support set_result operation')
195
196 def set_exception(self, exception):
197 raise RuntimeError('Task does not support set_exception operation')
198
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700199 def get_stack(self, *, limit=None):
200 """Return the list of stack frames for this task's coroutine.
201
Victor Stinnerd87de832014-12-02 17:57:04 +0100202 If the coroutine is not done, this returns the stack where it is
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700203 suspended. If the coroutine has completed successfully or was
204 cancelled, this returns an empty list. If the coroutine was
205 terminated by an exception, this returns the list of traceback
206 frames.
207
208 The frames are always ordered from oldest to newest.
209
Yury Selivanovb0b0e622014-02-18 22:27:48 -0500210 The optional limit gives the maximum number of frames to
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700211 return; by default all available frames are returned. Its
212 meaning differs depending on whether a stack or a traceback is
213 returned: the newest frames of a stack are returned, but the
214 oldest frames of a traceback are returned. (This matches the
215 behavior of the traceback module.)
216
217 For reasons beyond our control, only one stack frame is
218 returned for a suspended coroutine.
219 """
Yury Selivanova0c1ba62016-10-28 12:52:37 -0400220 return base_tasks._task_get_stack(self, limit)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700221
222 def print_stack(self, *, limit=None, file=None):
223 """Print the stack or traceback for this task's coroutine.
224
225 This produces output similar to that of the traceback module,
226 for the frames retrieved by get_stack(). The limit argument
227 is passed to get_stack(). The file argument is an I/O stream
R David Murray8e069d52014-09-24 13:13:45 -0400228 to which the output is written; by default output is written
229 to sys.stderr.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700230 """
Yury Selivanova0c1ba62016-10-28 12:52:37 -0400231 return base_tasks._task_print_stack(self, limit, file)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700232
233 def cancel(self):
R David Murray8e069d52014-09-24 13:13:45 -0400234 """Request that this task cancel itself.
Victor Stinner4bd652a2014-04-07 11:18:06 +0200235
Victor Stinner8d213572014-06-02 23:06:46 +0200236 This arranges for a CancelledError to be thrown into the
Victor Stinner4bd652a2014-04-07 11:18:06 +0200237 wrapped coroutine on the next cycle through the event loop.
238 The coroutine then has a chance to clean up or even deny
239 the request using try/except/finally.
240
R David Murray8e069d52014-09-24 13:13:45 -0400241 Unlike Future.cancel, this does not guarantee that the
Victor Stinner4bd652a2014-04-07 11:18:06 +0200242 task will be cancelled: the exception might be caught and
R David Murray8e069d52014-09-24 13:13:45 -0400243 acted upon, delaying cancellation of the task or preventing
244 cancellation completely. The task may also return a value or
245 raise a different exception.
Victor Stinner4bd652a2014-04-07 11:18:06 +0200246
247 Immediately after this method is called, Task.cancelled() will
248 not return True (unless the task was already cancelled). A
249 task will be marked as cancelled when the wrapped coroutine
250 terminates with a CancelledError exception (even if cancel()
251 was not called).
252 """
Yury Selivanov7ce1c6f2017-06-11 13:49:18 +0000253 self._log_traceback = False
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700254 if self.done():
255 return False
256 if self._fut_waiter is not None:
257 if self._fut_waiter.cancel():
258 # Leave self._fut_waiter; it may be a Task that
259 # catches and ignores the cancellation so we may have
260 # to cancel it again later.
261 return True
Yury Selivanov22feeb82018-01-24 11:31:01 -0500262 # It must be the case that self.__step is already scheduled.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700263 self._must_cancel = True
264 return True
265
Yury Selivanov22feeb82018-01-24 11:31:01 -0500266 def __step(self, exc=None):
Yury Selivanov0cf16f92017-12-25 10:48:15 -0500267 if self.done():
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700268 raise exceptions.InvalidStateError(
Yury Selivanov0cf16f92017-12-25 10:48:15 -0500269 f'_step(): already done: {self!r}, {exc!r}')
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700270 if self._must_cancel:
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700271 if not isinstance(exc, exceptions.CancelledError):
272 exc = exceptions.CancelledError()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700273 self._must_cancel = False
274 coro = self._coro
275 self._fut_waiter = None
Guido van Rossum1a605ed2013-12-06 12:57:40 -0800276
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200277 _enter_task(self._loop, self)
Yury Selivanovd59bba82015-11-20 12:41:03 -0500278 # Call either coro.throw(exc) or coro.send(None).
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700279 try:
Yury Selivanovd59bba82015-11-20 12:41:03 -0500280 if exc is None:
281 # We use the `send` method directly, because coroutines
282 # don't have `__iter__` and `__next__` methods.
283 result = coro.send(None)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700284 else:
Yury Selivanovd59bba82015-11-20 12:41:03 -0500285 result = coro.throw(exc)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700286 except StopIteration as exc:
INADA Naoki991adca2017-05-11 21:18:38 +0900287 if self._must_cancel:
288 # Task is cancelled right before coro stops.
289 self._must_cancel = False
Yury Selivanovedad4d82019-09-25 03:32:08 -0700290 super().cancel()
INADA Naoki991adca2017-05-11 21:18:38 +0900291 else:
Yury Selivanov0cf16f92017-12-25 10:48:15 -0500292 super().set_result(exc.value)
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700293 except exceptions.CancelledError:
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700294 super().cancel() # I.e., Future.cancel(self).
Yury Selivanov431b5402019-05-27 14:45:12 +0200295 except (KeyboardInterrupt, SystemExit) as exc:
Yury Selivanov0cf16f92017-12-25 10:48:15 -0500296 super().set_exception(exc)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700297 raise
Yury Selivanov431b5402019-05-27 14:45:12 +0200298 except BaseException as exc:
299 super().set_exception(exc)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700300 else:
Guido van Rossum1140a032016-09-09 12:54:54 -0700301 blocking = getattr(result, '_asyncio_future_blocking', None)
302 if blocking is not None:
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700303 # Yielded Future must come from Future.__iter__().
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500304 if futures._get_loop(result) is not self._loop:
Yury Selivanov6370f342017-12-10 18:36:12 -0500305 new_exc = RuntimeError(
306 f'Task {self!r} got Future '
307 f'{result!r} attached to a different loop')
Yury Selivanovf23746a2018-01-22 19:11:18 -0500308 self._loop.call_soon(
Yury Selivanov22feeb82018-01-24 11:31:01 -0500309 self.__step, new_exc, context=self._context)
Guido van Rossum1140a032016-09-09 12:54:54 -0700310 elif blocking:
Yury Selivanov4145c832016-10-09 12:19:12 -0400311 if result is self:
Yury Selivanov6370f342017-12-10 18:36:12 -0500312 new_exc = RuntimeError(
313 f'Task cannot await on itself: {self!r}')
Yury Selivanovf23746a2018-01-22 19:11:18 -0500314 self._loop.call_soon(
Yury Selivanov22feeb82018-01-24 11:31:01 -0500315 self.__step, new_exc, context=self._context)
Yury Selivanov4145c832016-10-09 12:19:12 -0400316 else:
317 result._asyncio_future_blocking = False
Yury Selivanovf23746a2018-01-22 19:11:18 -0500318 result.add_done_callback(
Yury Selivanov22feeb82018-01-24 11:31:01 -0500319 self.__wakeup, context=self._context)
Yury Selivanov4145c832016-10-09 12:19:12 -0400320 self._fut_waiter = result
321 if self._must_cancel:
322 if self._fut_waiter.cancel():
323 self._must_cancel = False
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700324 else:
Yury Selivanov6370f342017-12-10 18:36:12 -0500325 new_exc = RuntimeError(
326 f'yield was used instead of yield from '
327 f'in task {self!r} with {result!r}')
Yury Selivanovf23746a2018-01-22 19:11:18 -0500328 self._loop.call_soon(
Yury Selivanov22feeb82018-01-24 11:31:01 -0500329 self.__step, new_exc, context=self._context)
Yury Selivanov6370f342017-12-10 18:36:12 -0500330
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700331 elif result is None:
332 # Bare yield relinquishes control for one event loop iteration.
Yury Selivanov22feeb82018-01-24 11:31:01 -0500333 self._loop.call_soon(self.__step, context=self._context)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700334 elif inspect.isgenerator(result):
335 # Yielding a generator is just wrong.
Yury Selivanov6370f342017-12-10 18:36:12 -0500336 new_exc = RuntimeError(
337 f'yield was used instead of yield from for '
Serhiy Storchaka66553542018-05-20 16:30:31 +0300338 f'generator in task {self!r} with {result!r}')
Yury Selivanovf23746a2018-01-22 19:11:18 -0500339 self._loop.call_soon(
Yury Selivanov22feeb82018-01-24 11:31:01 -0500340 self.__step, new_exc, context=self._context)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700341 else:
342 # Yielding something else is an error.
Yury Selivanov6370f342017-12-10 18:36:12 -0500343 new_exc = RuntimeError(f'Task got bad yield: {result!r}')
Yury Selivanovf23746a2018-01-22 19:11:18 -0500344 self._loop.call_soon(
Yury Selivanov22feeb82018-01-24 11:31:01 -0500345 self.__step, new_exc, context=self._context)
Guido van Rossum1a605ed2013-12-06 12:57:40 -0800346 finally:
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200347 _leave_task(self._loop, self)
Victor Stinnerd74ac822014-03-04 23:07:08 +0100348 self = None # Needed to break cycles when an exception occurs.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700349
Yury Selivanov22feeb82018-01-24 11:31:01 -0500350 def __wakeup(self, future):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700351 try:
Yury Selivanova4afc482015-11-16 15:12:10 -0500352 future.result()
Yury Selivanov431b5402019-05-27 14:45:12 +0200353 except BaseException as exc:
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700354 # This may also be a cancellation.
Yury Selivanov22feeb82018-01-24 11:31:01 -0500355 self.__step(exc)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700356 else:
Yury Selivanova4afc482015-11-16 15:12:10 -0500357 # Don't pass the value of `future.result()` explicitly,
358 # as `Future.__iter__` and `Future.__await__` don't need it.
359 # If we call `_step(value, None)` instead of `_step()`,
360 # Python eval loop would use `.send(value)` method call,
361 # instead of `__next__()`, which is slower for futures
362 # that return non-generator iterators from their `__iter__`.
Yury Selivanov22feeb82018-01-24 11:31:01 -0500363 self.__step()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700364 self = None # Needed to break cycles when an exception occurs.
365
366
Yury Selivanova0c1ba62016-10-28 12:52:37 -0400367_PyTask = Task
368
369
370try:
371 import _asyncio
372except ImportError:
373 pass
374else:
375 # _CTask is needed for tests.
376 Task = _CTask = _asyncio.Task
377
378
Alex Grönholmcca4eec2018-08-09 00:06:47 +0300379def create_task(coro, *, name=None):
Andrew Svetlovf74ef452017-12-15 07:04:38 +0200380 """Schedule the execution of a coroutine object in a spawn task.
381
382 Return a Task object.
383 """
384 loop = events.get_running_loop()
Alex Grönholmcca4eec2018-08-09 00:06:47 +0300385 task = loop.create_task(coro)
386 _set_task_name(task, name)
387 return task
Andrew Svetlovf74ef452017-12-15 07:04:38 +0200388
389
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700390# wait() and as_completed() similar to those in PEP 3148.
391
392FIRST_COMPLETED = concurrent.futures.FIRST_COMPLETED
393FIRST_EXCEPTION = concurrent.futures.FIRST_EXCEPTION
394ALL_COMPLETED = concurrent.futures.ALL_COMPLETED
395
396
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200397async def wait(fs, *, loop=None, timeout=None, return_when=ALL_COMPLETED):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700398 """Wait for the Futures and coroutines given by fs to complete.
399
Victor Stinnerdb74d982014-06-10 11:16:05 +0200400 The sequence futures must not be empty.
401
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700402 Coroutines will be wrapped in Tasks.
403
404 Returns two sets of Future: (done, pending).
405
406 Usage:
407
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200408 done, pending = await asyncio.wait(fs)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700409
410 Note: This does not raise TimeoutError! Futures that aren't done
411 when the timeout occurs are returned in the second set.
412 """
Guido van Rossum7b3b3dc2016-09-09 14:26:31 -0700413 if futures.isfuture(fs) or coroutines.iscoroutine(fs):
Yury Selivanov6370f342017-12-10 18:36:12 -0500414 raise TypeError(f"expect a list of futures, not {type(fs).__name__}")
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700415 if not fs:
416 raise ValueError('Set of coroutines/Futures is empty.')
Victor Stinnere931f7b2014-07-16 18:50:39 +0200417 if return_when not in (FIRST_COMPLETED, FIRST_EXCEPTION, ALL_COMPLETED):
Yury Selivanov6370f342017-12-10 18:36:12 -0500418 raise ValueError(f'Invalid return_when value: {return_when}')
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700419
420 if loop is None:
João Júnior558c49b2018-09-24 06:51:22 -0300421 loop = events.get_running_loop()
422 else:
Matthias Bussonnierd0ebf132019-05-20 23:20:10 -0700423 warnings.warn("The loop argument is deprecated since Python 3.8, "
424 "and scheduled for removal in Python 3.10.",
João Júnior558c49b2018-09-24 06:51:22 -0300425 DeprecationWarning, stacklevel=2)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700426
Kyle Stanley89aa7f02019-12-30 06:50:19 -0500427 if any(coroutines.iscoroutine(f) for f in set(fs)):
428 warnings.warn("The explicit passing of coroutine objects to "
429 "asyncio.wait() is deprecated since Python 3.8, and "
430 "scheduled for removal in Python 3.11.",
431 DeprecationWarning, stacklevel=2)
432
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400433 fs = {ensure_future(f, loop=loop) for f in set(fs)}
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700434
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200435 return await _wait(fs, timeout, return_when, loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700436
437
Victor Stinner59e08022014-08-28 11:19:25 +0200438def _release_waiter(waiter, *args):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700439 if not waiter.done():
Victor Stinner59e08022014-08-28 11:19:25 +0200440 waiter.set_result(None)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700441
442
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200443async def wait_for(fut, timeout, *, loop=None):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700444 """Wait for the single Future or coroutine to complete, with timeout.
445
446 Coroutine will be wrapped in Task.
447
Victor Stinner421e49b2014-01-23 17:40:59 +0100448 Returns result of the Future or coroutine. When a timeout occurs,
449 it cancels the task and raises TimeoutError. To avoid the task
450 cancellation, wrap it in shield().
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700451
Victor Stinner922bc2c2015-01-15 16:29:10 +0100452 If the wait is cancelled, the task is also cancelled.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700453
Victor Stinner922bc2c2015-01-15 16:29:10 +0100454 This function is a coroutine.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700455 """
456 if loop is None:
João Júnior558c49b2018-09-24 06:51:22 -0300457 loop = events.get_running_loop()
458 else:
Matthias Bussonnierd0ebf132019-05-20 23:20:10 -0700459 warnings.warn("The loop argument is deprecated since Python 3.8, "
460 "and scheduled for removal in Python 3.10.",
João Júnior558c49b2018-09-24 06:51:22 -0300461 DeprecationWarning, stacklevel=2)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700462
Guido van Rossum48c66c32014-01-29 14:30:38 -0800463 if timeout is None:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200464 return await fut
Guido van Rossum48c66c32014-01-29 14:30:38 -0800465
Victor K4d071892017-10-05 19:04:39 +0300466 if timeout <= 0:
467 fut = ensure_future(fut, loop=loop)
468
469 if fut.done():
470 return fut.result()
471
472 fut.cancel()
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700473 raise exceptions.TimeoutError()
Victor K4d071892017-10-05 19:04:39 +0300474
Yury Selivanov7661db62016-05-16 15:38:39 -0400475 waiter = loop.create_future()
Victor Stinner59e08022014-08-28 11:19:25 +0200476 timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
477 cb = functools.partial(_release_waiter, waiter)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700478
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400479 fut = ensure_future(fut, loop=loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700480 fut.add_done_callback(cb)
481
482 try:
Victor Stinner59e08022014-08-28 11:19:25 +0200483 # wait until the future completes or the timeout
Victor Stinner922bc2c2015-01-15 16:29:10 +0100484 try:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200485 await waiter
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700486 except exceptions.CancelledError:
Victor Stinner922bc2c2015-01-15 16:29:10 +0100487 fut.remove_done_callback(cb)
488 fut.cancel()
489 raise
Victor Stinner59e08022014-08-28 11:19:25 +0200490
491 if fut.done():
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700492 return fut.result()
493 else:
494 fut.remove_done_callback(cb)
Elvis Pranskevichuse2b340a2018-05-29 17:31:01 -0400495 # We must ensure that the task is not running
496 # after wait_for() returns.
497 # See https://bugs.python.org/issue32751
498 await _cancel_and_wait(fut, loop=loop)
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700499 raise exceptions.TimeoutError()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700500 finally:
501 timeout_handle.cancel()
502
503
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200504async def _wait(fs, timeout, return_when, loop):
Elvis Pranskevichuse2b340a2018-05-29 17:31:01 -0400505 """Internal helper for wait().
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700506
507 The fs argument must be a collection of Futures.
508 """
509 assert fs, 'Set of Futures is empty.'
Yury Selivanov7661db62016-05-16 15:38:39 -0400510 waiter = loop.create_future()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700511 timeout_handle = None
512 if timeout is not None:
513 timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
514 counter = len(fs)
515
516 def _on_completion(f):
517 nonlocal counter
518 counter -= 1
519 if (counter <= 0 or
520 return_when == FIRST_COMPLETED or
521 return_when == FIRST_EXCEPTION and (not f.cancelled() and
522 f.exception() is not None)):
523 if timeout_handle is not None:
524 timeout_handle.cancel()
525 if not waiter.done():
Victor Stinner59e08022014-08-28 11:19:25 +0200526 waiter.set_result(None)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700527
528 for f in fs:
529 f.add_done_callback(_on_completion)
530
531 try:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200532 await waiter
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700533 finally:
534 if timeout_handle is not None:
535 timeout_handle.cancel()
gescheitc1964e92019-05-03 18:18:02 +0300536 for f in fs:
537 f.remove_done_callback(_on_completion)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700538
539 done, pending = set(), set()
540 for f in fs:
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700541 if f.done():
542 done.add(f)
543 else:
544 pending.add(f)
545 return done, pending
546
547
Elvis Pranskevichuse2b340a2018-05-29 17:31:01 -0400548async def _cancel_and_wait(fut, loop):
549 """Cancel the *fut* future or task and wait until it completes."""
550
551 waiter = loop.create_future()
552 cb = functools.partial(_release_waiter, waiter)
553 fut.add_done_callback(cb)
554
555 try:
556 fut.cancel()
557 # We cannot wait on *fut* directly to make
558 # sure _cancel_and_wait itself is reliably cancellable.
559 await waiter
560 finally:
561 fut.remove_done_callback(cb)
562
563
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700564# This is *not* a @coroutine! It is just an iterator (yielding Futures).
565def as_completed(fs, *, loop=None, timeout=None):
Guido van Rossumb58f0532014-02-12 17:58:19 -0800566 """Return an iterator whose values are coroutines.
567
568 When waiting for the yielded coroutines you'll get the results (or
569 exceptions!) of the original Futures (or coroutines), in the order
570 in which and as soon as they complete.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700571
572 This differs from PEP 3148; the proper way to use this is:
573
574 for f in as_completed(fs):
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200575 result = await f # The 'await' may raise.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700576 # Use result.
577
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200578 If a timeout is specified, the 'await' will raise
Guido van Rossumb58f0532014-02-12 17:58:19 -0800579 TimeoutError when the timeout occurs before all Futures are done.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700580
581 Note: The futures 'f' are not necessarily members of fs.
582 """
Guido van Rossum7b3b3dc2016-09-09 14:26:31 -0700583 if futures.isfuture(fs) or coroutines.iscoroutine(fs):
Yury Selivanov6370f342017-12-10 18:36:12 -0500584 raise TypeError(f"expect a list of futures, not {type(fs).__name__}")
Andrew Svetlova4888792019-09-12 15:40:40 +0300585
Guido van Rossumb58f0532014-02-12 17:58:19 -0800586 from .queues import Queue # Import here to avoid circular import problem.
587 done = Queue(loop=loop)
Andrew Svetlova4888792019-09-12 15:40:40 +0300588
589 if loop is None:
590 loop = events.get_event_loop()
591 else:
592 warnings.warn("The loop argument is deprecated since Python 3.8, "
593 "and scheduled for removal in Python 3.10.",
594 DeprecationWarning, stacklevel=2)
595 todo = {ensure_future(f, loop=loop) for f in set(fs)}
Guido van Rossumb58f0532014-02-12 17:58:19 -0800596 timeout_handle = None
597
598 def _on_timeout():
599 for f in todo:
600 f.remove_done_callback(_on_completion)
601 done.put_nowait(None) # Queue a dummy value for _wait_for_one().
602 todo.clear() # Can't do todo.remove(f) in the loop.
603
604 def _on_completion(f):
605 if not todo:
606 return # _on_timeout() was here first.
607 todo.remove(f)
608 done.put_nowait(f)
609 if not todo and timeout_handle is not None:
610 timeout_handle.cancel()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700611
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200612 async def _wait_for_one():
613 f = await done.get()
Guido van Rossumb58f0532014-02-12 17:58:19 -0800614 if f is None:
615 # Dummy value from _on_timeout().
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700616 raise exceptions.TimeoutError
Guido van Rossumb58f0532014-02-12 17:58:19 -0800617 return f.result() # May raise f.exception().
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700618
Guido van Rossumb58f0532014-02-12 17:58:19 -0800619 for f in todo:
620 f.add_done_callback(_on_completion)
621 if todo and timeout is not None:
622 timeout_handle = loop.call_later(timeout, _on_timeout)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700623 for _ in range(len(todo)):
624 yield _wait_for_one()
625
626
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200627@types.coroutine
628def __sleep0():
629 """Skip one event loop run cycle.
630
631 This is a private helper for 'asyncio.sleep()', used
632 when the 'delay' is set to 0. It uses a bare 'yield'
Yury Selivanov22feeb82018-01-24 11:31:01 -0500633 expression (which Task.__step knows how to handle)
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200634 instead of creating a Future object.
635 """
636 yield
637
638
639async def sleep(delay, result=None, *, loop=None):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700640 """Coroutine that completes after a given time (in seconds)."""
Andrew Svetlov5382c052017-12-17 16:41:30 +0200641 if delay <= 0:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200642 await __sleep0()
Yury Selivanovade04122015-11-05 14:29:04 -0500643 return result
644
Yury Selivanov7661db62016-05-16 15:38:39 -0400645 if loop is None:
João Júnior558c49b2018-09-24 06:51:22 -0300646 loop = events.get_running_loop()
647 else:
Matthias Bussonnierd0ebf132019-05-20 23:20:10 -0700648 warnings.warn("The loop argument is deprecated since Python 3.8, "
649 "and scheduled for removal in Python 3.10.",
João Júnior558c49b2018-09-24 06:51:22 -0300650 DeprecationWarning, stacklevel=2)
651
Yury Selivanov7661db62016-05-16 15:38:39 -0400652 future = loop.create_future()
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500653 h = loop.call_later(delay,
654 futures._set_result_unless_cancelled,
655 future, result)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700656 try:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200657 return await future
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700658 finally:
659 h.cancel()
660
661
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400662def ensure_future(coro_or_future, *, loop=None):
Yury Selivanov620279b2015-10-02 15:00:19 -0400663 """Wrap a coroutine or an awaitable in a future.
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400664
665 If the argument is a Future, it is returned directly.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700666 """
jimmylaie549c4b2018-05-28 06:42:05 -1000667 if coroutines.iscoroutine(coro_or_future):
Victor Stinner896a25a2014-07-08 11:29:25 +0200668 if loop is None:
669 loop = events.get_event_loop()
670 task = loop.create_task(coro_or_future)
Victor Stinner80f53aa2014-06-27 13:52:20 +0200671 if task._source_traceback:
672 del task._source_traceback[-1]
673 return task
jimmylaie549c4b2018-05-28 06:42:05 -1000674 elif futures.isfuture(coro_or_future):
675 if loop is not None and loop is not futures._get_loop(coro_or_future):
Zackery Spytz4737b922019-05-03 09:35:26 -0600676 raise ValueError('The future belongs to a different loop than '
677 'the one specified as the loop argument')
jimmylaie549c4b2018-05-28 06:42:05 -1000678 return coro_or_future
Victor Stinner3f438a92017-11-28 14:43:52 +0100679 elif inspect.isawaitable(coro_or_future):
Yury Selivanov620279b2015-10-02 15:00:19 -0400680 return ensure_future(_wrap_awaitable(coro_or_future), loop=loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700681 else:
Charles Renwickae5b3262017-04-21 16:49:48 -0400682 raise TypeError('An asyncio.Future, a coroutine or an awaitable is '
683 'required')
Yury Selivanov620279b2015-10-02 15:00:19 -0400684
685
Andrew Svetlov68b34a72019-05-16 17:52:10 +0300686@types.coroutine
Yury Selivanov620279b2015-10-02 15:00:19 -0400687def _wrap_awaitable(awaitable):
688 """Helper for asyncio.ensure_future().
689
690 Wraps awaitable (an object with __await__) into a coroutine
691 that will later be wrapped in a Task by ensure_future().
692 """
693 return (yield from awaitable.__await__())
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700694
Andrew Svetlov68b34a72019-05-16 17:52:10 +0300695_wrap_awaitable._is_coroutine = _is_coroutine
696
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700697
698class _GatheringFuture(futures.Future):
699 """Helper for gather().
700
701 This overrides cancel() to cancel all the children and act more
702 like Task.cancel(), which doesn't immediately mark itself as
703 cancelled.
704 """
705
706 def __init__(self, children, *, loop=None):
707 super().__init__(loop=loop)
708 self._children = children
Yury Selivanov863b6742018-05-29 17:20:02 -0400709 self._cancel_requested = False
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700710
711 def cancel(self):
712 if self.done():
713 return False
Yury Selivanov3d676152016-10-21 17:22:17 -0400714 ret = False
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700715 for child in self._children:
Yury Selivanov3d676152016-10-21 17:22:17 -0400716 if child.cancel():
717 ret = True
Yury Selivanov863b6742018-05-29 17:20:02 -0400718 if ret:
719 # If any child tasks were actually cancelled, we should
720 # propagate the cancellation request regardless of
721 # *return_exceptions* argument. See issue 32684.
722 self._cancel_requested = True
Yury Selivanov3d676152016-10-21 17:22:17 -0400723 return ret
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700724
725
726def gather(*coros_or_futures, loop=None, return_exceptions=False):
Yury Selivanov36c2c042017-12-19 07:19:53 -0500727 """Return a future aggregating results from the given coroutines/futures.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700728
Guido van Rossume3c65a72016-09-30 08:17:15 -0700729 Coroutines will be wrapped in a future and scheduled in the event
730 loop. They will not necessarily be scheduled in the same order as
731 passed in.
732
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700733 All futures must share the same event loop. If all the tasks are
734 done successfully, the returned future's result is the list of
735 results (in the order of the original sequence, not necessarily
Yury Selivanovf317cb72014-02-06 12:03:53 -0500736 the order of results arrival). If *return_exceptions* is True,
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700737 exceptions in the tasks are treated the same as successful
738 results, and gathered in the result list; otherwise, the first
739 raised exception will be immediately propagated to the returned
740 future.
741
742 Cancellation: if the outer Future is cancelled, all children (that
743 have not completed yet) are also cancelled. If any child is
744 cancelled, this is treated as if it raised CancelledError --
745 the outer Future is *not* cancelled in this case. (This is to
746 prevent the cancellation of one child to cause other children to
747 be cancelled.)
748 """
Victor Stinnerf03b3c72014-07-16 18:36:24 +0200749 if not coros_or_futures:
Yury Selivanov7661db62016-05-16 15:38:39 -0400750 if loop is None:
751 loop = events.get_event_loop()
Andrew Svetlova4888792019-09-12 15:40:40 +0300752 else:
753 warnings.warn("The loop argument is deprecated since Python 3.8, "
754 "and scheduled for removal in Python 3.10.",
755 DeprecationWarning, stacklevel=2)
Yury Selivanov7661db62016-05-16 15:38:39 -0400756 outer = loop.create_future()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700757 outer.set_result([])
758 return outer
Victor Stinnerf03b3c72014-07-16 18:36:24 +0200759
Yury Selivanov36c2c042017-12-19 07:19:53 -0500760 def _done_callback(fut):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700761 nonlocal nfinished
Yury Selivanov36c2c042017-12-19 07:19:53 -0500762 nfinished += 1
763
Victor Stinner3531d902015-01-09 01:42:52 +0100764 if outer.done():
765 if not fut.cancelled():
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700766 # Mark exception retrieved.
767 fut.exception()
768 return
Victor Stinner3531d902015-01-09 01:42:52 +0100769
Yury Selivanov36c2c042017-12-19 07:19:53 -0500770 if not return_exceptions:
771 if fut.cancelled():
772 # Check if 'fut' is cancelled first, as
773 # 'fut.exception()' will *raise* a CancelledError
774 # instead of returning it.
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700775 exc = exceptions.CancelledError()
Yury Selivanov36c2c042017-12-19 07:19:53 -0500776 outer.set_exception(exc)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700777 return
Yury Selivanov36c2c042017-12-19 07:19:53 -0500778 else:
779 exc = fut.exception()
780 if exc is not None:
781 outer.set_exception(exc)
782 return
783
784 if nfinished == nfuts:
785 # All futures are done; create a list of results
786 # and set it to the 'outer' future.
787 results = []
788
789 for fut in children:
790 if fut.cancelled():
791 # Check if 'fut' is cancelled first, as
792 # 'fut.exception()' will *raise* a CancelledError
793 # instead of returning it.
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700794 res = exceptions.CancelledError()
Yury Selivanov36c2c042017-12-19 07:19:53 -0500795 else:
796 res = fut.exception()
797 if res is None:
798 res = fut.result()
799 results.append(res)
800
Yury Selivanov863b6742018-05-29 17:20:02 -0400801 if outer._cancel_requested:
802 # If gather is being cancelled we must propagate the
803 # cancellation regardless of *return_exceptions* argument.
804 # See issue 32684.
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700805 outer.set_exception(exceptions.CancelledError())
Yury Selivanov863b6742018-05-29 17:20:02 -0400806 else:
807 outer.set_result(results)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700808
Yury Selivanov36c2c042017-12-19 07:19:53 -0500809 arg_to_fut = {}
810 children = []
811 nfuts = 0
812 nfinished = 0
813 for arg in coros_or_futures:
814 if arg not in arg_to_fut:
815 fut = ensure_future(arg, loop=loop)
816 if loop is None:
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500817 loop = futures._get_loop(fut)
Yury Selivanov36c2c042017-12-19 07:19:53 -0500818 if fut is not arg:
819 # 'arg' was not a Future, therefore, 'fut' is a new
820 # Future created specifically for 'arg'. Since the caller
821 # can't control it, disable the "destroy pending task"
822 # warning.
823 fut._log_destroy_pending = False
824
825 nfuts += 1
826 arg_to_fut[arg] = fut
827 fut.add_done_callback(_done_callback)
828
829 else:
830 # There's a duplicate Future object in coros_or_futures.
831 fut = arg_to_fut[arg]
832
833 children.append(fut)
834
835 outer = _GatheringFuture(children, loop=loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700836 return outer
837
838
839def shield(arg, *, loop=None):
840 """Wait for a future, shielding it from cancellation.
841
842 The statement
843
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200844 res = await shield(something())
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700845
846 is exactly equivalent to the statement
847
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200848 res = await something()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700849
850 *except* that if the coroutine containing it is cancelled, the
851 task running in something() is not cancelled. From the POV of
852 something(), the cancellation did not happen. But its caller is
853 still cancelled, so the yield-from expression still raises
854 CancelledError. Note: If something() is cancelled by other means
855 this will still cancel shield().
856
857 If you want to completely ignore cancellation (not recommended)
858 you can combine shield() with a try/except clause, as follows:
859
860 try:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200861 res = await shield(something())
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700862 except CancelledError:
863 res = None
864 """
Andrew Svetlova4888792019-09-12 15:40:40 +0300865 if loop is not None:
866 warnings.warn("The loop argument is deprecated since Python 3.8, "
867 "and scheduled for removal in Python 3.10.",
868 DeprecationWarning, stacklevel=2)
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400869 inner = ensure_future(arg, loop=loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700870 if inner.done():
871 # Shortcut.
872 return inner
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500873 loop = futures._get_loop(inner)
Yury Selivanov7661db62016-05-16 15:38:39 -0400874 outer = loop.create_future()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700875
Romain Picardb35acc52019-05-07 20:58:24 +0200876 def _inner_done_callback(inner):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700877 if outer.cancelled():
Victor Stinner3531d902015-01-09 01:42:52 +0100878 if not inner.cancelled():
879 # Mark inner's result as retrieved.
880 inner.exception()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700881 return
Victor Stinner3531d902015-01-09 01:42:52 +0100882
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700883 if inner.cancelled():
884 outer.cancel()
885 else:
886 exc = inner.exception()
887 if exc is not None:
888 outer.set_exception(exc)
889 else:
890 outer.set_result(inner.result())
891
Romain Picardb35acc52019-05-07 20:58:24 +0200892
893 def _outer_done_callback(outer):
894 if not inner.done():
895 inner.remove_done_callback(_inner_done_callback)
896
897 inner.add_done_callback(_inner_done_callback)
898 outer.add_done_callback(_outer_done_callback)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700899 return outer
Guido van Rossum841d9ee2015-10-03 08:31:42 -0700900
901
902def run_coroutine_threadsafe(coro, loop):
903 """Submit a coroutine object to a given event loop.
904
905 Return a concurrent.futures.Future to access the result.
906 """
907 if not coroutines.iscoroutine(coro):
908 raise TypeError('A coroutine object is required')
909 future = concurrent.futures.Future()
910
911 def callback():
Guido van Rossum601953b2015-10-05 16:20:00 -0700912 try:
913 futures._chain_future(ensure_future(coro, loop=loop), future)
Yury Selivanov431b5402019-05-27 14:45:12 +0200914 except (SystemExit, KeyboardInterrupt):
915 raise
916 except BaseException as exc:
Guido van Rossum601953b2015-10-05 16:20:00 -0700917 if future.set_running_or_notify_cancel():
918 future.set_exception(exc)
919 raise
Guido van Rossum841d9ee2015-10-03 08:31:42 -0700920
921 loop.call_soon_threadsafe(callback)
922 return future
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200923
924
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500925# WeakSet containing all alive tasks.
926_all_tasks = weakref.WeakSet()
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200927
928# Dictionary containing tasks that are currently active in
929# all running event loops. {EventLoop: Task}
930_current_tasks = {}
931
932
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500933def _register_task(task):
934 """Register a new task in asyncio as executed by loop."""
935 _all_tasks.add(task)
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200936
937
938def _enter_task(loop, task):
939 current_task = _current_tasks.get(loop)
940 if current_task is not None:
941 raise RuntimeError(f"Cannot enter into task {task!r} while another "
942 f"task {current_task!r} is being executed.")
943 _current_tasks[loop] = task
944
945
946def _leave_task(loop, task):
947 current_task = _current_tasks.get(loop)
948 if current_task is not task:
949 raise RuntimeError(f"Leaving task {task!r} does not match "
950 f"the current task {current_task!r}.")
951 del _current_tasks[loop]
952
953
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500954def _unregister_task(task):
955 """Unregister a task."""
956 _all_tasks.discard(task)
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200957
958
959_py_register_task = _register_task
960_py_unregister_task = _unregister_task
961_py_enter_task = _enter_task
962_py_leave_task = _leave_task
963
964
965try:
966 from _asyncio import (_register_task, _unregister_task,
967 _enter_task, _leave_task,
968 _all_tasks, _current_tasks)
969except ImportError:
970 pass
971else:
972 _c_register_task = _register_task
973 _c_unregister_task = _unregister_task
974 _c_enter_task = _enter_task
975 _c_leave_task = _leave_task