blob: 894d28eb107acd7b439478152c4a78c2264f8f12 [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
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400427 fs = {ensure_future(f, loop=loop) for f in set(fs)}
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700428
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200429 return await _wait(fs, timeout, return_when, loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700430
431
Victor Stinner59e08022014-08-28 11:19:25 +0200432def _release_waiter(waiter, *args):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700433 if not waiter.done():
Victor Stinner59e08022014-08-28 11:19:25 +0200434 waiter.set_result(None)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700435
436
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200437async def wait_for(fut, timeout, *, loop=None):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700438 """Wait for the single Future or coroutine to complete, with timeout.
439
440 Coroutine will be wrapped in Task.
441
Victor Stinner421e49b2014-01-23 17:40:59 +0100442 Returns result of the Future or coroutine. When a timeout occurs,
443 it cancels the task and raises TimeoutError. To avoid the task
444 cancellation, wrap it in shield().
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700445
Victor Stinner922bc2c2015-01-15 16:29:10 +0100446 If the wait is cancelled, the task is also cancelled.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700447
Victor Stinner922bc2c2015-01-15 16:29:10 +0100448 This function is a coroutine.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700449 """
450 if loop is None:
João Júnior558c49b2018-09-24 06:51:22 -0300451 loop = events.get_running_loop()
452 else:
Matthias Bussonnierd0ebf132019-05-20 23:20:10 -0700453 warnings.warn("The loop argument is deprecated since Python 3.8, "
454 "and scheduled for removal in Python 3.10.",
João Júnior558c49b2018-09-24 06:51:22 -0300455 DeprecationWarning, stacklevel=2)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700456
Guido van Rossum48c66c32014-01-29 14:30:38 -0800457 if timeout is None:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200458 return await fut
Guido van Rossum48c66c32014-01-29 14:30:38 -0800459
Victor K4d071892017-10-05 19:04:39 +0300460 if timeout <= 0:
461 fut = ensure_future(fut, loop=loop)
462
463 if fut.done():
464 return fut.result()
465
466 fut.cancel()
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700467 raise exceptions.TimeoutError()
Victor K4d071892017-10-05 19:04:39 +0300468
Yury Selivanov7661db62016-05-16 15:38:39 -0400469 waiter = loop.create_future()
Victor Stinner59e08022014-08-28 11:19:25 +0200470 timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
471 cb = functools.partial(_release_waiter, waiter)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700472
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400473 fut = ensure_future(fut, loop=loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700474 fut.add_done_callback(cb)
475
476 try:
Victor Stinner59e08022014-08-28 11:19:25 +0200477 # wait until the future completes or the timeout
Victor Stinner922bc2c2015-01-15 16:29:10 +0100478 try:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200479 await waiter
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700480 except exceptions.CancelledError:
Victor Stinner922bc2c2015-01-15 16:29:10 +0100481 fut.remove_done_callback(cb)
482 fut.cancel()
483 raise
Victor Stinner59e08022014-08-28 11:19:25 +0200484
485 if fut.done():
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700486 return fut.result()
487 else:
488 fut.remove_done_callback(cb)
Elvis Pranskevichuse2b340a2018-05-29 17:31:01 -0400489 # We must ensure that the task is not running
490 # after wait_for() returns.
491 # See https://bugs.python.org/issue32751
492 await _cancel_and_wait(fut, loop=loop)
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700493 raise exceptions.TimeoutError()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700494 finally:
495 timeout_handle.cancel()
496
497
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200498async def _wait(fs, timeout, return_when, loop):
Elvis Pranskevichuse2b340a2018-05-29 17:31:01 -0400499 """Internal helper for wait().
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700500
501 The fs argument must be a collection of Futures.
502 """
503 assert fs, 'Set of Futures is empty.'
Yury Selivanov7661db62016-05-16 15:38:39 -0400504 waiter = loop.create_future()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700505 timeout_handle = None
506 if timeout is not None:
507 timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
508 counter = len(fs)
509
510 def _on_completion(f):
511 nonlocal counter
512 counter -= 1
513 if (counter <= 0 or
514 return_when == FIRST_COMPLETED or
515 return_when == FIRST_EXCEPTION and (not f.cancelled() and
516 f.exception() is not None)):
517 if timeout_handle is not None:
518 timeout_handle.cancel()
519 if not waiter.done():
Victor Stinner59e08022014-08-28 11:19:25 +0200520 waiter.set_result(None)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700521
522 for f in fs:
523 f.add_done_callback(_on_completion)
524
525 try:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200526 await waiter
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700527 finally:
528 if timeout_handle is not None:
529 timeout_handle.cancel()
gescheitc1964e92019-05-03 18:18:02 +0300530 for f in fs:
531 f.remove_done_callback(_on_completion)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700532
533 done, pending = set(), set()
534 for f in fs:
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700535 if f.done():
536 done.add(f)
537 else:
538 pending.add(f)
539 return done, pending
540
541
Elvis Pranskevichuse2b340a2018-05-29 17:31:01 -0400542async def _cancel_and_wait(fut, loop):
543 """Cancel the *fut* future or task and wait until it completes."""
544
545 waiter = loop.create_future()
546 cb = functools.partial(_release_waiter, waiter)
547 fut.add_done_callback(cb)
548
549 try:
550 fut.cancel()
551 # We cannot wait on *fut* directly to make
552 # sure _cancel_and_wait itself is reliably cancellable.
553 await waiter
554 finally:
555 fut.remove_done_callback(cb)
556
557
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700558# This is *not* a @coroutine! It is just an iterator (yielding Futures).
559def as_completed(fs, *, loop=None, timeout=None):
Guido van Rossumb58f0532014-02-12 17:58:19 -0800560 """Return an iterator whose values are coroutines.
561
562 When waiting for the yielded coroutines you'll get the results (or
563 exceptions!) of the original Futures (or coroutines), in the order
564 in which and as soon as they complete.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700565
566 This differs from PEP 3148; the proper way to use this is:
567
568 for f in as_completed(fs):
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200569 result = await f # The 'await' may raise.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700570 # Use result.
571
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200572 If a timeout is specified, the 'await' will raise
Guido van Rossumb58f0532014-02-12 17:58:19 -0800573 TimeoutError when the timeout occurs before all Futures are done.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700574
575 Note: The futures 'f' are not necessarily members of fs.
576 """
Guido van Rossum7b3b3dc2016-09-09 14:26:31 -0700577 if futures.isfuture(fs) or coroutines.iscoroutine(fs):
Yury Selivanov6370f342017-12-10 18:36:12 -0500578 raise TypeError(f"expect a list of futures, not {type(fs).__name__}")
Andrew Svetlova4888792019-09-12 15:40:40 +0300579
Guido van Rossumb58f0532014-02-12 17:58:19 -0800580 from .queues import Queue # Import here to avoid circular import problem.
581 done = Queue(loop=loop)
Andrew Svetlova4888792019-09-12 15:40:40 +0300582
583 if loop is None:
584 loop = events.get_event_loop()
585 else:
586 warnings.warn("The loop argument is deprecated since Python 3.8, "
587 "and scheduled for removal in Python 3.10.",
588 DeprecationWarning, stacklevel=2)
589 todo = {ensure_future(f, loop=loop) for f in set(fs)}
Guido van Rossumb58f0532014-02-12 17:58:19 -0800590 timeout_handle = None
591
592 def _on_timeout():
593 for f in todo:
594 f.remove_done_callback(_on_completion)
595 done.put_nowait(None) # Queue a dummy value for _wait_for_one().
596 todo.clear() # Can't do todo.remove(f) in the loop.
597
598 def _on_completion(f):
599 if not todo:
600 return # _on_timeout() was here first.
601 todo.remove(f)
602 done.put_nowait(f)
603 if not todo and timeout_handle is not None:
604 timeout_handle.cancel()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700605
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200606 async def _wait_for_one():
607 f = await done.get()
Guido van Rossumb58f0532014-02-12 17:58:19 -0800608 if f is None:
609 # Dummy value from _on_timeout().
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700610 raise exceptions.TimeoutError
Guido van Rossumb58f0532014-02-12 17:58:19 -0800611 return f.result() # May raise f.exception().
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700612
Guido van Rossumb58f0532014-02-12 17:58:19 -0800613 for f in todo:
614 f.add_done_callback(_on_completion)
615 if todo and timeout is not None:
616 timeout_handle = loop.call_later(timeout, _on_timeout)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700617 for _ in range(len(todo)):
618 yield _wait_for_one()
619
620
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200621@types.coroutine
622def __sleep0():
623 """Skip one event loop run cycle.
624
625 This is a private helper for 'asyncio.sleep()', used
626 when the 'delay' is set to 0. It uses a bare 'yield'
Yury Selivanov22feeb82018-01-24 11:31:01 -0500627 expression (which Task.__step knows how to handle)
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200628 instead of creating a Future object.
629 """
630 yield
631
632
633async def sleep(delay, result=None, *, loop=None):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700634 """Coroutine that completes after a given time (in seconds)."""
Andrew Svetlov5382c052017-12-17 16:41:30 +0200635 if delay <= 0:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200636 await __sleep0()
Yury Selivanovade04122015-11-05 14:29:04 -0500637 return result
638
Yury Selivanov7661db62016-05-16 15:38:39 -0400639 if loop is None:
João Júnior558c49b2018-09-24 06:51:22 -0300640 loop = events.get_running_loop()
641 else:
Matthias Bussonnierd0ebf132019-05-20 23:20:10 -0700642 warnings.warn("The loop argument is deprecated since Python 3.8, "
643 "and scheduled for removal in Python 3.10.",
João Júnior558c49b2018-09-24 06:51:22 -0300644 DeprecationWarning, stacklevel=2)
645
Yury Selivanov7661db62016-05-16 15:38:39 -0400646 future = loop.create_future()
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500647 h = loop.call_later(delay,
648 futures._set_result_unless_cancelled,
649 future, result)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700650 try:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200651 return await future
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700652 finally:
653 h.cancel()
654
655
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400656def ensure_future(coro_or_future, *, loop=None):
Yury Selivanov620279b2015-10-02 15:00:19 -0400657 """Wrap a coroutine or an awaitable in a future.
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400658
659 If the argument is a Future, it is returned directly.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700660 """
jimmylaie549c4b2018-05-28 06:42:05 -1000661 if coroutines.iscoroutine(coro_or_future):
Victor Stinner896a25a2014-07-08 11:29:25 +0200662 if loop is None:
663 loop = events.get_event_loop()
664 task = loop.create_task(coro_or_future)
Victor Stinner80f53aa2014-06-27 13:52:20 +0200665 if task._source_traceback:
666 del task._source_traceback[-1]
667 return task
jimmylaie549c4b2018-05-28 06:42:05 -1000668 elif futures.isfuture(coro_or_future):
669 if loop is not None and loop is not futures._get_loop(coro_or_future):
Zackery Spytz4737b922019-05-03 09:35:26 -0600670 raise ValueError('The future belongs to a different loop than '
671 'the one specified as the loop argument')
jimmylaie549c4b2018-05-28 06:42:05 -1000672 return coro_or_future
Victor Stinner3f438a92017-11-28 14:43:52 +0100673 elif inspect.isawaitable(coro_or_future):
Yury Selivanov620279b2015-10-02 15:00:19 -0400674 return ensure_future(_wrap_awaitable(coro_or_future), loop=loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700675 else:
Charles Renwickae5b3262017-04-21 16:49:48 -0400676 raise TypeError('An asyncio.Future, a coroutine or an awaitable is '
677 'required')
Yury Selivanov620279b2015-10-02 15:00:19 -0400678
679
Andrew Svetlov68b34a72019-05-16 17:52:10 +0300680@types.coroutine
Yury Selivanov620279b2015-10-02 15:00:19 -0400681def _wrap_awaitable(awaitable):
682 """Helper for asyncio.ensure_future().
683
684 Wraps awaitable (an object with __await__) into a coroutine
685 that will later be wrapped in a Task by ensure_future().
686 """
687 return (yield from awaitable.__await__())
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700688
Andrew Svetlov68b34a72019-05-16 17:52:10 +0300689_wrap_awaitable._is_coroutine = _is_coroutine
690
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700691
692class _GatheringFuture(futures.Future):
693 """Helper for gather().
694
695 This overrides cancel() to cancel all the children and act more
696 like Task.cancel(), which doesn't immediately mark itself as
697 cancelled.
698 """
699
700 def __init__(self, children, *, loop=None):
701 super().__init__(loop=loop)
702 self._children = children
Yury Selivanov863b6742018-05-29 17:20:02 -0400703 self._cancel_requested = False
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700704
705 def cancel(self):
706 if self.done():
707 return False
Yury Selivanov3d676152016-10-21 17:22:17 -0400708 ret = False
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700709 for child in self._children:
Yury Selivanov3d676152016-10-21 17:22:17 -0400710 if child.cancel():
711 ret = True
Yury Selivanov863b6742018-05-29 17:20:02 -0400712 if ret:
713 # If any child tasks were actually cancelled, we should
714 # propagate the cancellation request regardless of
715 # *return_exceptions* argument. See issue 32684.
716 self._cancel_requested = True
Yury Selivanov3d676152016-10-21 17:22:17 -0400717 return ret
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700718
719
720def gather(*coros_or_futures, loop=None, return_exceptions=False):
Yury Selivanov36c2c042017-12-19 07:19:53 -0500721 """Return a future aggregating results from the given coroutines/futures.
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700722
Guido van Rossume3c65a72016-09-30 08:17:15 -0700723 Coroutines will be wrapped in a future and scheduled in the event
724 loop. They will not necessarily be scheduled in the same order as
725 passed in.
726
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700727 All futures must share the same event loop. If all the tasks are
728 done successfully, the returned future's result is the list of
729 results (in the order of the original sequence, not necessarily
Yury Selivanovf317cb72014-02-06 12:03:53 -0500730 the order of results arrival). If *return_exceptions* is True,
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700731 exceptions in the tasks are treated the same as successful
732 results, and gathered in the result list; otherwise, the first
733 raised exception will be immediately propagated to the returned
734 future.
735
736 Cancellation: if the outer Future is cancelled, all children (that
737 have not completed yet) are also cancelled. If any child is
738 cancelled, this is treated as if it raised CancelledError --
739 the outer Future is *not* cancelled in this case. (This is to
740 prevent the cancellation of one child to cause other children to
741 be cancelled.)
742 """
Victor Stinnerf03b3c72014-07-16 18:36:24 +0200743 if not coros_or_futures:
Yury Selivanov7661db62016-05-16 15:38:39 -0400744 if loop is None:
745 loop = events.get_event_loop()
Andrew Svetlova4888792019-09-12 15:40:40 +0300746 else:
747 warnings.warn("The loop argument is deprecated since Python 3.8, "
748 "and scheduled for removal in Python 3.10.",
749 DeprecationWarning, stacklevel=2)
Yury Selivanov7661db62016-05-16 15:38:39 -0400750 outer = loop.create_future()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700751 outer.set_result([])
752 return outer
Victor Stinnerf03b3c72014-07-16 18:36:24 +0200753
Yury Selivanov36c2c042017-12-19 07:19:53 -0500754 def _done_callback(fut):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700755 nonlocal nfinished
Yury Selivanov36c2c042017-12-19 07:19:53 -0500756 nfinished += 1
757
Victor Stinner3531d902015-01-09 01:42:52 +0100758 if outer.done():
759 if not fut.cancelled():
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700760 # Mark exception retrieved.
761 fut.exception()
762 return
Victor Stinner3531d902015-01-09 01:42:52 +0100763
Yury Selivanov36c2c042017-12-19 07:19:53 -0500764 if not return_exceptions:
765 if fut.cancelled():
766 # Check if 'fut' is cancelled first, as
767 # 'fut.exception()' will *raise* a CancelledError
768 # instead of returning it.
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700769 exc = exceptions.CancelledError()
Yury Selivanov36c2c042017-12-19 07:19:53 -0500770 outer.set_exception(exc)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700771 return
Yury Selivanov36c2c042017-12-19 07:19:53 -0500772 else:
773 exc = fut.exception()
774 if exc is not None:
775 outer.set_exception(exc)
776 return
777
778 if nfinished == nfuts:
779 # All futures are done; create a list of results
780 # and set it to the 'outer' future.
781 results = []
782
783 for fut in children:
784 if fut.cancelled():
785 # Check if 'fut' is cancelled first, as
786 # 'fut.exception()' will *raise* a CancelledError
787 # instead of returning it.
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700788 res = exceptions.CancelledError()
Yury Selivanov36c2c042017-12-19 07:19:53 -0500789 else:
790 res = fut.exception()
791 if res is None:
792 res = fut.result()
793 results.append(res)
794
Yury Selivanov863b6742018-05-29 17:20:02 -0400795 if outer._cancel_requested:
796 # If gather is being cancelled we must propagate the
797 # cancellation regardless of *return_exceptions* argument.
798 # See issue 32684.
Andrew Svetlov0baa72f2018-09-11 10:13:04 -0700799 outer.set_exception(exceptions.CancelledError())
Yury Selivanov863b6742018-05-29 17:20:02 -0400800 else:
801 outer.set_result(results)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700802
Yury Selivanov36c2c042017-12-19 07:19:53 -0500803 arg_to_fut = {}
804 children = []
805 nfuts = 0
806 nfinished = 0
807 for arg in coros_or_futures:
808 if arg not in arg_to_fut:
809 fut = ensure_future(arg, loop=loop)
810 if loop is None:
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500811 loop = futures._get_loop(fut)
Yury Selivanov36c2c042017-12-19 07:19:53 -0500812 if fut is not arg:
813 # 'arg' was not a Future, therefore, 'fut' is a new
814 # Future created specifically for 'arg'. Since the caller
815 # can't control it, disable the "destroy pending task"
816 # warning.
817 fut._log_destroy_pending = False
818
819 nfuts += 1
820 arg_to_fut[arg] = fut
821 fut.add_done_callback(_done_callback)
822
823 else:
824 # There's a duplicate Future object in coros_or_futures.
825 fut = arg_to_fut[arg]
826
827 children.append(fut)
828
829 outer = _GatheringFuture(children, loop=loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700830 return outer
831
832
833def shield(arg, *, loop=None):
834 """Wait for a future, shielding it from cancellation.
835
836 The statement
837
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200838 res = await shield(something())
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700839
840 is exactly equivalent to the statement
841
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200842 res = await something()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700843
844 *except* that if the coroutine containing it is cancelled, the
845 task running in something() is not cancelled. From the POV of
846 something(), the cancellation did not happen. But its caller is
847 still cancelled, so the yield-from expression still raises
848 CancelledError. Note: If something() is cancelled by other means
849 this will still cancel shield().
850
851 If you want to completely ignore cancellation (not recommended)
852 you can combine shield() with a try/except clause, as follows:
853
854 try:
Andrew Svetlov5f841b52017-12-09 00:23:48 +0200855 res = await shield(something())
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700856 except CancelledError:
857 res = None
858 """
Andrew Svetlova4888792019-09-12 15:40:40 +0300859 if loop is not None:
860 warnings.warn("The loop argument is deprecated since Python 3.8, "
861 "and scheduled for removal in Python 3.10.",
862 DeprecationWarning, stacklevel=2)
Yury Selivanov59eb9a42015-05-11 14:48:38 -0400863 inner = ensure_future(arg, loop=loop)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700864 if inner.done():
865 # Shortcut.
866 return inner
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500867 loop = futures._get_loop(inner)
Yury Selivanov7661db62016-05-16 15:38:39 -0400868 outer = loop.create_future()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700869
Romain Picardb35acc52019-05-07 20:58:24 +0200870 def _inner_done_callback(inner):
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700871 if outer.cancelled():
Victor Stinner3531d902015-01-09 01:42:52 +0100872 if not inner.cancelled():
873 # Mark inner's result as retrieved.
874 inner.exception()
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700875 return
Victor Stinner3531d902015-01-09 01:42:52 +0100876
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700877 if inner.cancelled():
878 outer.cancel()
879 else:
880 exc = inner.exception()
881 if exc is not None:
882 outer.set_exception(exc)
883 else:
884 outer.set_result(inner.result())
885
Romain Picardb35acc52019-05-07 20:58:24 +0200886
887 def _outer_done_callback(outer):
888 if not inner.done():
889 inner.remove_done_callback(_inner_done_callback)
890
891 inner.add_done_callback(_inner_done_callback)
892 outer.add_done_callback(_outer_done_callback)
Guido van Rossum27b7c7e2013-10-17 13:40:50 -0700893 return outer
Guido van Rossum841d9ee2015-10-03 08:31:42 -0700894
895
896def run_coroutine_threadsafe(coro, loop):
897 """Submit a coroutine object to a given event loop.
898
899 Return a concurrent.futures.Future to access the result.
900 """
901 if not coroutines.iscoroutine(coro):
902 raise TypeError('A coroutine object is required')
903 future = concurrent.futures.Future()
904
905 def callback():
Guido van Rossum601953b2015-10-05 16:20:00 -0700906 try:
907 futures._chain_future(ensure_future(coro, loop=loop), future)
Yury Selivanov431b5402019-05-27 14:45:12 +0200908 except (SystemExit, KeyboardInterrupt):
909 raise
910 except BaseException as exc:
Guido van Rossum601953b2015-10-05 16:20:00 -0700911 if future.set_running_or_notify_cancel():
912 future.set_exception(exc)
913 raise
Guido van Rossum841d9ee2015-10-03 08:31:42 -0700914
915 loop.call_soon_threadsafe(callback)
916 return future
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200917
918
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500919# WeakSet containing all alive tasks.
920_all_tasks = weakref.WeakSet()
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200921
922# Dictionary containing tasks that are currently active in
923# all running event loops. {EventLoop: Task}
924_current_tasks = {}
925
926
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500927def _register_task(task):
928 """Register a new task in asyncio as executed by loop."""
929 _all_tasks.add(task)
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200930
931
932def _enter_task(loop, task):
933 current_task = _current_tasks.get(loop)
934 if current_task is not None:
935 raise RuntimeError(f"Cannot enter into task {task!r} while another "
936 f"task {current_task!r} is being executed.")
937 _current_tasks[loop] = task
938
939
940def _leave_task(loop, task):
941 current_task = _current_tasks.get(loop)
942 if current_task is not task:
943 raise RuntimeError(f"Leaving task {task!r} does not match "
944 f"the current task {current_task!r}.")
945 del _current_tasks[loop]
946
947
Yury Selivanovca9b36c2017-12-23 15:04:15 -0500948def _unregister_task(task):
949 """Unregister a task."""
950 _all_tasks.discard(task)
Andrew Svetlov44d1a592017-12-16 21:58:38 +0200951
952
953_py_register_task = _register_task
954_py_unregister_task = _unregister_task
955_py_enter_task = _enter_task
956_py_leave_task = _leave_task
957
958
959try:
960 from _asyncio import (_register_task, _unregister_task,
961 _enter_task, _leave_task,
962 _all_tasks, _current_tasks)
963except ImportError:
964 pass
965else:
966 _c_register_task = _register_task
967 _c_unregister_task = _unregister_task
968 _c_enter_task = _enter_task
969 _c_leave_task = _leave_task