__all__ = (

'Task', 'create_task',

'FIRST_COMPLETED', 'FIRST_EXCEPTION', 'ALL_COMPLETED',

'wait', 'wait_for', 'as_completed', 'sleep',

'gather', 'shield', 'ensure_future', 'run_coroutine_threadsafe',

'current_task', 'all_tasks',

'create_eager_task_factory', 'eager_task_factory',

'_register_task', '_unregister_task', '_enter_task', '_leave_task',

)

import concurrent.futures

import contextvars

import functools

import inspect

import itertools

import math

import types

import weakref

from types import GenericAlias

from . import base_tasks

from . import coroutines

from . import events

from . import exceptions

from . import futures

from . import queues

from . import timeouts

_task_name_counter = itertools.count(1).__next__

def current_task(loop=None):

"""Return a currently executed task."""

if loop is None:

loop = events.get_running_loop()

return _current_tasks.get(loop)

def all_tasks(loop=None):

"""Return a set of all tasks for the loop."""

if loop is None:

loop = events.get_running_loop()

# capturing the set of eager tasks first, so if an eager task "graduates"

# to a regular task in another thread, we don't risk missing it.

eager_tasks = list(_eager_tasks)

return {t for t in itertools.chain(_scheduled_tasks, eager_tasks)

if futures._get_loop(t) is loop and not t.done()}

class Task(futures._PyFuture): # Inherit Python Task implementation

# from a Python Future implementation.

"""A coroutine wrapped in a Future."""

# An important invariant maintained while a Task not done:

# _fut_waiter is either None or a Future. The Future

# can be either done() or not done().

# The task can be in any of 3 states:

#

# - 1: _fut_waiter is not None and not _fut_waiter.done():

# __step() is *not* scheduled and the Task is waiting for _fut_waiter.

# - 2: (_fut_waiter is None or _fut_waiter.done()) and __step() is scheduled:

# the Task is waiting for __step() to be executed.

# - 3: _fut_waiter is None and __step() is *not* scheduled:

# the Task is currently executing (in __step()).

#

# * In state 1, one of the callbacks of __fut_waiter must be __wakeup().

# * The transition from 1 to 2 happens when _fut_waiter becomes done(),

# as it schedules __wakeup() to be called (which calls __step() so

# we way that __step() is scheduled).

# * It transitions from 2 to 3 when __step() is executed, and it clears

# _fut_waiter to None.

# If False, don't log a message if the task is destroyed while its

# status is still pending

_log_destroy_pending = True

def __init__(self, coro, *, loop=None, name=None, context=None,

eager_start=False):

super().__init__(loop=loop)

if self._source_traceback:

del self._source_traceback[-1]

if not coroutines.iscoroutine(coro):

# raise after Future.__init__(), attrs are required for __del__

# prevent logging for pending task in __del__

self._log_destroy_pending = False

raise TypeError(f"a coroutine was expected, got {coro!r}")

if name is None:

self._name = f'Task-{_task_name_counter()}'

else:

self._name = str(name)

self._num_cancels_requested = 0

self._must_cancel = False

self._fut_waiter = None

self._coro = coro

if context is None:

self._context = contextvars.copy_context()

else:

self._context = context

if eager_start and self._loop.is_running():

self.__eager_start()

else:

self._loop.call_soon(self.__step, context=self._context)

_py_register_task(self)

def __del__(self):

if self._state == futures._PENDING and self._log_destroy_pending:

context = {

'task': self,

'message': 'Task was destroyed but it is pending!',

}

if self._source_traceback:

context['source_traceback'] = self._source_traceback

self._loop.call_exception_handler(context)

super().__del__()

__class_getitem__ = classmethod(GenericAlias)

def __repr__(self):

return base_tasks._task_repr(self)

def get_coro(self):

return self._coro

def get_context(self):

return self._context

def get_name(self):

return self._name

def set_name(self, value):

self._name = str(value)

def set_result(self, result):

raise RuntimeError('Task does not support set_result operation')

def set_exception(self, exception):

raise RuntimeError('Task does not support set_exception operation')

def get_stack(self, *, limit=None):

"""Return the list of stack frames for this task's coroutine.

return base_tasks._task_get_stack(self, limit)

def print_stack(self, *, limit=None, file=None):

"""Print the stack or traceback for this task's coroutine.

return base_tasks._task_print_stack(self, limit, file)

def cancel(self, msg=None):

"""Request that this task cancel itself.

self._log_traceback = False

if self.done():

return False

self._num_cancels_requested += 1

# These two lines are controversial. See discussion starting at

# https://github.com/python/cpython/pull/31394#issuecomment-1053545331

# Also remember that this is duplicated in _asynciomodule.c.

# if self._num_cancels_requested > 1:

# return False

if self._fut_waiter is not None:

if self._fut_waiter.cancel(msg=msg):

# Leave self._fut_waiter; it may be a Task that

# catches and ignores the cancellation so we may have

# to cancel it again later.

return True

# It must be the case that self.__step is already scheduled.

self._must_cancel = True

self._cancel_message = msg

return True

def cancelling(self):

"""Return the count of the task's cancellation requests.

return self._num_cancels_requested

def uncancel(self):

"""Decrement the task's count of cancellation requests.

if self._num_cancels_requested > 0:

self._num_cancels_requested -= 1

if self._num_cancels_requested == 0:

self._must_cancel = False

return self._num_cancels_requested

def __eager_start(self):

prev_task = _py_swap_current_task(self._loop, self)

try:

_py_register_eager_task(self)

try:

self._context.run(self.__step_run_and_handle_result, None)

finally:

_py_unregister_eager_task(self)

finally:

try:

curtask = _py_swap_current_task(self._loop, prev_task)

assert curtask is self

finally:

if self.done():

self._coro = None

self = None # Needed to break cycles when an exception occurs.

else:

_py_register_task(self)

def __step(self, exc=None):

if self.done():

raise exceptions.InvalidStateError(

f'__step(): already done: {self!r}, {exc!r}')

if self._must_cancel:

if not isinstance(exc, exceptions.CancelledError):

exc = self._make_cancelled_error()

self._must_cancel = False

self._fut_waiter = None

_py_enter_task(self._loop, self)

try:

self.__step_run_and_handle_result(exc)

finally:

_py_leave_task(self._loop, self)

self = None # Needed to break cycles when an exception occurs.

def __step_run_and_handle_result(self, exc):

coro = self._coro

try:

if exc is None:

# We use the `send` method directly, because coroutines

# don't have `__iter__` and `__next__` methods.

result = coro.send(None)

else:

result = coro.throw(exc)

except StopIteration as exc:

if self._must_cancel:

# Task is cancelled right before coro stops.

self._must_cancel = False

super().cancel(msg=self._cancel_message)

else:

super().set_result(exc.value)

except exceptions.CancelledError as exc:

# Save the original exception so we can chain it later.

self._cancelled_exc = exc

super().cancel() # I.e., Future.cancel(self).

except (KeyboardInterrupt, SystemExit) as exc:

super().set_exception(exc)

raise

except BaseException as exc:

super().set_exception(exc)

else:

blocking = getattr(result, '_asyncio_future_blocking', None)

if blocking is not None:

# Yielded Future must come from Future.__iter__().

if futures._get_loop(result) is not self._loop:

new_exc = RuntimeError(

f'Task {self!r} got Future '

f'{result!r} attached to a different loop')

self._loop.call_soon(

self.__step, new_exc, context=self._context)

elif blocking:

if result is self:

new_exc = RuntimeError(

f'Task cannot await on itself: {self!r}')

self._loop.call_soon(

self.__step, new_exc, context=self._context)

else:

futures.future_add_to_awaited_by(result, self)

result._asyncio_future_blocking = False

result.add_done_callback(

self.__wakeup, context=self._context)

self._fut_waiter = result

if self._must_cancel:

if self._fut_waiter.cancel(

msg=self._cancel_message):

self._must_cancel = False

else:

new_exc = RuntimeError(

f'yield was used instead of yield from '

f'in task {self!r} with {result!r}')

self._loop.call_soon(

self.__step, new_exc, context=self._context)

elif result is None:

# Bare yield relinquishes control for one event loop iteration.

self._loop.call_soon(self.__step, context=self._context)

elif inspect.isgenerator(result):

# Yielding a generator is just wrong.

new_exc = RuntimeError(

f'yield was used instead of yield from for '

f'generator in task {self!r} with {result!r}')

self._loop.call_soon(

self.__step, new_exc, context=self._context)

else:

# Yielding something else is an error.

new_exc = RuntimeError(f'Task got bad yield: {result!r}')

self._loop.call_soon(

self.__step, new_exc, context=self._context)

finally:

self = None # Needed to break cycles when an exception occurs.

def __wakeup(self, future):

futures.future_discard_from_awaited_by(future, self)

try:

future.result()

except BaseException as exc:

# This may also be a cancellation.

self.__step(exc)

else:

# Don't pass the value of `future.result()` explicitly,

# as `Future.__iter__` and `Future.__await__` don't need it.

# If we call `__step(value, None)` instead of `__step()`,

# Python eval loop would use `.send(value)` method call,

# instead of `__next__()`, which is slower for futures

# that return non-generator iterators from their `__iter__`.

self.__step()

self = None # Needed to break cycles when an exception occurs.

_PyTask = Task

try:

import _asyncio

except ImportError:

pass

else:

# _CTask is needed for tests.

Task = _CTask = _asyncio.Task

def create_task(coro, **kwargs):

"""Schedule the execution of a coroutine object in a spawn task.

loop = events.get_running_loop()

return loop.create_task(coro, **kwargs)

FIRST_COMPLETED = concurrent.futures.FIRST_COMPLETED

FIRST_EXCEPTION = concurrent.futures.FIRST_EXCEPTION

ALL_COMPLETED = concurrent.futures.ALL_COMPLETED

async def wait(fs, *, timeout=None, return_when=ALL_COMPLETED):

"""Wait for the Futures or Tasks given by fs to complete.

if futures.isfuture(fs) or coroutines.iscoroutine(fs):

raise TypeError(f"expect a list of futures, not {type(fs).__name__}")

if not fs:

raise ValueError('Set of Tasks/Futures is empty.')

if return_when not in (FIRST_COMPLETED, FIRST_EXCEPTION, ALL_COMPLETED):

raise ValueError(f'Invalid return_when value: {return_when}')

fs = set(fs)

if any(coroutines.iscoroutine(f) for f in fs):

raise TypeError("Passing coroutines is forbidden, use tasks explicitly.")

loop = events.get_running_loop()

return await _wait(fs, timeout, return_when, loop)

def _release_waiter(waiter, *args):

if not waiter.done():

waiter.set_result(None)

async def wait_for(fut, timeout):

"""Wait for the single Future or coroutine to complete, with timeout.

# The special case for timeout <= 0 is for the following case:

#

# async def test_waitfor():

# func_started = False

#

# async def func():

# nonlocal func_started

# func_started = True

#

# try:

# await asyncio.wait_for(func(), 0)

# except asyncio.TimeoutError:

# assert not func_started

# else:

# assert False

#

# asyncio.run(test_waitfor())

if timeout is not None and timeout <= 0:

fut = ensure_future(fut)

if fut.done():

return fut.result()

await _cancel_and_wait(fut)

try:

return fut.result()

except exceptions.CancelledError as exc:

raise TimeoutError from exc

async with timeouts.timeout(timeout):

return await fut

async def _wait(fs, timeout, return_when, loop):

"""Internal helper for wait().

assert fs, 'Set of Futures is empty.'

waiter = loop.create_future()

timeout_handle = None

if timeout is not None:

timeout_handle = loop.call_later(timeout, _release_waiter, waiter)

counter = len(fs)

cur_task = current_task()

def _on_completion(f):

nonlocal counter

counter -= 1

if (counter <= 0 or

return_when == FIRST_COMPLETED or

return_when == FIRST_EXCEPTION and (not f.cancelled() and

f.exception() is not None)):

if timeout_handle is not None:

timeout_handle.cancel()

if not waiter.done():

waiter.set_result(None)

futures.future_discard_from_awaited_by(f, cur_task)

for f in fs:

f.add_done_callback(_on_completion)

futures.future_add_to_awaited_by(f, cur_task)

try:

await waiter

finally:

if timeout_handle is not None:

timeout_handle.cancel()

for f in fs:

f.remove_done_callback(_on_completion)

futures.future_discard_from_awaited_by(f, cur_task)

done, pending = set(), set()

for f in fs:

if f.done():

done.add(f)

else:

pending.add(f)

return done, pending

async def _cancel_and_wait(fut):

"""Cancel the *fut* future or task and wait until it completes."""

loop = events.get_running_loop()

waiter = loop.create_future()

cb = functools.partial(_release_waiter, waiter)

fut.add_done_callback(cb)

# gh-157058: awaiting the waiter leaves no edge on fut, add it here

cur_task = current_task()

futures.future_add_to_awaited_by(fut, cur_task)

try:

fut.cancel()

# We cannot wait on *fut* directly to make

# sure _cancel_and_wait itself is reliably cancellable.

await waiter

finally:

fut.remove_done_callback(cb)

futures.future_discard_from_awaited_by(fut, cur_task)

class _AsCompletedIterator:

"""Iterator of awaitables representing tasks of asyncio.as_completed.

def __init__(self, aws, timeout):

self._done = queues.Queue()

self._timeout_handle = None

loop = events.get_event_loop()

self._cur_task = current_task()

todo = {ensure_future(aw, loop=loop) for aw in set(aws)}

for f in todo:

f.add_done_callback(self._handle_completion)

futures.future_add_to_awaited_by(f, self._cur_task)

if todo and timeout is not None:

self._timeout_handle = (

loop.call_later(timeout, self._handle_timeout)

)

self._todo = todo

self._todo_left = len(todo)

def __aiter__(self):

return self

def __iter__(self):

return self

async def __anext__(self):

if not self._todo_left:

raise StopAsyncIteration

assert self._todo_left > 0

self._todo_left -= 1

return await self._wait_for_one()

def __next__(self):

if not self._todo_left:

raise StopIteration

assert self._todo_left > 0

self._todo_left -= 1

return self._wait_for_one(resolve=True)

def _handle_timeout(self):

for f in self._todo:

f.remove_done_callback(self._handle_completion)

futures.future_discard_from_awaited_by(f, self._cur_task)

self._done.put_nowait(None) # Sentinel for _wait_for_one().

self._todo.clear() # Can't do todo.remove(f) in the loop.

def _handle_completion(self, f):

if not self._todo:

return # _handle_timeout() was here first.

self._todo.remove(f)

futures.future_discard_from_awaited_by(f, self._cur_task)

self._done.put_nowait(f)

if not self._todo and self._timeout_handle is not None:

self._timeout_handle.cancel()

async def _wait_for_one(self, resolve=False):

# Wait for the next future to be done and return it unless resolve is

# set, in which case return either the result of the future or raise

# an exception.

f = await self._done.get()

if f is None:

# Dummy value from _handle_timeout().

raise exceptions.TimeoutError

return f.result() if resolve else f

def as_completed(fs, *, timeout=None):

"""Create an iterator of awaitables or their results in completion order.

if inspect.isawaitable(fs):

raise TypeError(

f"expects an iterable of awaitables, not {type(fs).__name__}"

)

return _AsCompletedIterator(fs, timeout)

def __sleep0():

"""Skip one event loop run cycle.

yield

async def sleep(delay, result=None):

"""Coroutine that completes after a given time (in seconds)."""

if delay <= 0:

await __sleep0()

return result

if math.isnan(delay):

raise ValueError("Invalid delay: NaN (not a number)")

loop = events.get_running_loop()

future = loop.create_future()

h = loop.call_later(delay,

futures._set_result_unless_cancelled,

future, result)

try:

return await future

finally:

h.cancel()

def ensure_future(coro_or_future, *, loop=None):

"""Wrap a coroutine or an awaitable in a future.

if futures.isfuture(coro_or_future):

if loop is not None and loop is not futures._get_loop(coro_or_future):

raise ValueError('The future belongs to a different loop than '

'the one specified as the loop argument')

return coro_or_future

should_close = True

if not coroutines.iscoroutine(coro_or_future):

if inspect.isawaitable(coro_or_future):

async def _wrap_awaitable(awaitable):

return await awaitable

coro_or_future = _wrap_awaitable(coro_or_future)

should_close = False

else:

raise TypeError('An asyncio.Future, a coroutine or an awaitable '

'is required')

if loop is None:

loop = events.get_event_loop()

try:

return loop.create_task(coro_or_future)

except RuntimeError:

if should_close:

coro_or_future.close()

raise

class _GatheringFuture(futures.Future):

"""Helper for gather().

def __init__(self, children, *, loop):

assert loop is not None

super().__init__(loop=loop)

self._children = children

self._cancel_requested = False

def cancel(self, msg=None):

if self.done():

return False

ret = False

for child in self._children:

if child.cancel(msg=msg):

ret = True

if ret:

# If any child tasks were actually cancelled, we should

# propagate the cancellation request regardless of

# *return_exceptions* argument. See issue 32684.

self._cancel_requested = True

return ret

def _discard_awaited_by(children, waiter, outer):

for fut in children:

futures.future_discard_from_awaited_by(fut, waiter)

def gather(*coros_or_futures, return_exceptions=False):

"""Return a future aggregating results from the given coroutines/futures.

if not coros_or_futures:

loop = events.get_event_loop()

outer = loop.create_future()

outer.set_result([])

return outer

loop = events._get_running_loop()

if loop is not None:

cur_task = current_task(loop)

else:

cur_task = None

def _done_callback(fut, cur_task=cur_task):

nonlocal nfinished

nfinished += 1

if cur_task is not None:

futures.future_discard_from_awaited_by(fut, cur_task)

if outer is None or outer.done():

if not fut.cancelled():

# Mark exception retrieved.

fut.exception()

return

if not return_exceptions:

if fut.cancelled():

# Check if 'fut' is cancelled first, as

# 'fut.exception()' will *raise* a CancelledError

# instead of returning it.

exc = fut._make_cancelled_error()

outer.set_exception(exc)

return

else:

exc = fut.exception()

if exc is not None:

outer.set_exception(exc)

return

if nfinished == nfuts:

# All futures are done; create a list of results

# and set it to the 'outer' future.

results = []

for fut in children:

if fut.cancelled():

# Check if 'fut' is cancelled first, as 'fut.exception()'

# will *raise* a CancelledError instead of returning it.

# Also, since we're adding the exception return value

# to 'results' instead of raising it, don't bother

# setting __context__. This also lets us preserve

# calling '_make_cancelled_error()' at most once.

res = exceptions.CancelledError(

'' if fut._cancel_message is None else

fut._cancel_message)

else:

res = fut.exception()

if res is None:

res = fut.result()

results.append(res)

if outer._cancel_requested:

# If gather is being cancelled we must propagate the

# cancellation regardless of *return_exceptions* argument.

# See issue 32684.

exc = fut._make_cancelled_error()

outer.set_exception(exc)

else:

outer.set_result(results)

arg_to_fut = {}

children = []

nfuts = 0

nfinished = 0

done_futs = []

outer = None # bpo-46672

for arg in coros_or_futures:

if arg not in arg_to_fut:

fut = ensure_future(arg, loop=loop)

if loop is None:

loop = futures._get_loop(fut)

if fut is not arg:

# 'arg' was not a Future, therefore, 'fut' is a new

# Future created specifically for 'arg'. Since the caller

# can't control it, disable the "destroy pending task"

# warning.

fut._log_destroy_pending = False

nfuts += 1

arg_to_fut[arg] = fut

if fut.done():

done_futs.append(fut)

else:

if cur_task is not None:

futures.future_add_to_awaited_by(fut, cur_task)

fut.add_done_callback(_done_callback)

else:

# There's a duplicate Future object in coros_or_futures.

fut = arg_to_fut[arg]

children.append(fut)

outer = _GatheringFuture(children, loop=loop)

if cur_task is not None:

# gh-157213: a child outliving gather() must lose the awaited-by edge

outer.add_done_callback(

functools.partial(_discard_awaited_by, children, cur_task))

# Run done callbacks after GatheringFuture created so any post-processing

# can be performed at this point

# optimization: in the special case that *all* futures finished eagerly,

# this will effectively complete the gather eagerly, with the last

# callback setting the result (or exception) on outer before returning it

for fut in done_futs:

_done_callback(fut)

return outer

def shield(arg):

"""Wait for a future, shielding it from cancellation.

inner = ensure_future(arg)

if inner.done():

# Shortcut.

return inner

loop = futures._get_loop(inner)

outer = loop.create_future()

if loop is not None and (cur_task := current_task(loop)) is not None:

futures.future_add_to_awaited_by(inner, cur_task)

else:

cur_task = None

def _clear_awaited_by_callback(inner):

futures.future_discard_from_awaited_by(inner, cur_task)

def _inner_done_callback(inner):

if outer.cancelled():

return

if inner.cancelled():

outer.cancel()

else:

exc = inner.exception()

if exc is not None:

outer.set_exception(exc)

else:

outer.set_result(inner.result())

def _outer_done_callback(outer):

if not inner.done():

inner.remove_done_callback(_inner_done_callback)

if cur_task is not None:

inner.remove_done_callback(_clear_awaited_by_callback)

futures.future_discard_from_awaited_by(inner, cur_task)

if cur_task is not None:

inner.add_done_callback(_clear_awaited_by_callback)