diff --git a/nest_asyncio.py b/nest_asyncio.py index 01ec2ab..70d1618 100644 --- a/nest_asyncio.py +++ b/nest_asyncio.py @@ -13,7 +13,6 @@ def apply(loop=None): """Patch asyncio to make its event loop reentrant.""" _patch_asyncio() _patch_policy() - _patch_task() _patch_tornado() loop = loop or asyncio.get_event_loop() @@ -126,9 +125,20 @@ def _run_once(self): break handle = ready.popleft() if not handle._cancelled: - handle._run() + # preempt the current task so that that checks in + # Task.__step do not raise + curr_task = curr_tasks.pop(self, None) + + try: + handle._run() + finally: + # restore the current task + if curr_task is not None: + curr_tasks[self] = curr_task + handle = None + @contextmanager def manage_run(self): """Set up the loop for running.""" @@ -193,45 +203,11 @@ def _check_running(self): os.name == 'nt' and issubclass(cls, asyncio.ProactorEventLoop)) if sys.version_info < (3, 7, 0): cls._set_coroutine_origin_tracking = cls._set_coroutine_wrapper + curr_tasks = asyncio.tasks._current_tasks \ + if sys.version_info >= (3, 7, 0) else asyncio.Task._current_tasks cls._nest_patched = True -def _patch_task(): - """Patch the Task's step and enter/leave methods to make it reentrant.""" - - def step(task, exc=None): - curr_task = curr_tasks.get(task._loop) - try: - step_orig(task, exc) - finally: - if curr_task is None: - curr_tasks.pop(task._loop, None) - else: - curr_tasks[task._loop] = curr_task - - Task = asyncio.Task - if hasattr(Task, '_nest_patched'): - return - if sys.version_info >= (3, 7, 0): - - def enter_task(loop, task): - curr_tasks[loop] = task - - def leave_task(loop, task): - curr_tasks.pop(loop, None) - - asyncio.tasks._enter_task = enter_task - asyncio.tasks._leave_task = leave_task - curr_tasks = asyncio.tasks._current_tasks - step_orig = Task._Task__step - Task._Task__step = step - else: - curr_tasks = Task._current_tasks - step_orig = Task._step - Task._step = step - Task._nest_patched = True - - def _patch_tornado(): """ If tornado is imported before nest_asyncio, make tornado aware of