diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py index 11858a0274a69f..4ee6b69d109f01 100644 --- a/Lib/asyncio/futures.py +++ b/Lib/asyncio/futures.py @@ -72,6 +72,8 @@ class Future: __log_traceback = False + _eager_result = False + def __init__(self, *, loop=None): """Initialize the future. @@ -168,6 +170,30 @@ def cancel(self, msg=None): self.__schedule_callbacks() return True + def __execute_callbacks(self): + callbacks = self._callbacks[:] + if not callbacks: + return + + self._callbacks[:] = [] + for callback, ctx in callbacks: + # TODO: run_in_context + try: + callback(self) + except BaseException as exc: + cb = format_helpers._format_callback_source( + callback, (self,), + debug=self._loop.get_debug()) + msg = f'Exception in callback {cb}' + context = { + 'message': msg, + 'exception': exc, + 'handle': self, + } + if self._source_traceback: + context['source_traceback'] = self._source_traceback + self._loop.call_exception_handler(context) + def __schedule_callbacks(self): """Internal: Ask the event loop to call all callbacks. @@ -270,6 +296,19 @@ def set_result(self, result): self._state = _FINISHED self.__schedule_callbacks() + def eager_set_result(self, result): + """Mark the future done and set its result. + + If the future is already done when this method is called, raises + InvalidStateError. + """ + if self._state != _PENDING: + raise exceptions.InvalidStateError(f'{self._state}: {self!r}') + self._result = result + self._state = _FINISHED + self._eager_result = True + self.__execute_callbacks() + def set_exception(self, exception): """Mark the future done and set an exception. diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index fbd5c39a7c56ac..474295a2d92a03 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -370,7 +370,10 @@ def __wakeup(self, future): # 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() + if future._eager_result: + self.__eager_start() + else: + self.__step() self = None # Needed to break cycles when an exception occurs. diff --git a/eager_future.py b/eager_future.py new file mode 100644 index 00000000000000..8e11aa7b9e5da4 --- /dev/null +++ b/eager_future.py @@ -0,0 +1,42 @@ +from functools import partial, Placeholder as _ + +import asyncio + +print(asyncio.Future) # Will show if it's the C or Python version +print(asyncio.Task) + +def callback_no_exc(fut, idx): + print(f"callback_no_throw: {idx}") + + +def callback_with_exc(fut): + print("callback_with_exception") + raise RuntimeError("TEST EXCEPTION") + + +async def func(fut): + try: + print("func: entered") + await fut + print("func: done") + except Exception as exc: + print(f"func: EXCEPTION: {exc}") + + +async def main(): + loop = asyncio.get_running_loop() + loop.set_task_factory(asyncio.eager_task_factory) + fut = asyncio.Future(loop=loop) + fut.add_done_callback(partial(callback_no_exc, _, 1)) + fut.add_done_callback(callback_with_exc) + fut.add_done_callback(partial(callback_no_exc, _, 2)) + loop.create_task(func(fut)) + await asyncio.sleep(1) + print("main: before set_result") + fut.eager_set_result(None) + print("main: set_result done") + await asyncio.sleep(1) + print("main: done") + + +asyncio.run(main())