231 lines
7.4 KiB
Python
231 lines
7.4 KiB
Python
import asyncio
|
|
import asyncio.events as events
|
|
import os
|
|
import sys
|
|
import threading
|
|
from heapq import heappop
|
|
|
|
|
|
def apply(loop=None):
|
|
"""Patch asyncio to make its event loop reentrent."""
|
|
loop = loop or asyncio.get_event_loop()
|
|
if not isinstance(loop, asyncio.BaseEventLoop):
|
|
raise ValueError('Can\'t patch loop of type %s' % type(loop))
|
|
if getattr(loop, '_nest_patched', None):
|
|
# already patched
|
|
return
|
|
_patch_asyncio()
|
|
_patch_loop(loop)
|
|
_patch_task()
|
|
_patch_handle()
|
|
_patch_tornado()
|
|
|
|
|
|
def _patch_asyncio():
|
|
"""
|
|
Patch asyncio module to use pure Python tasks and futures,
|
|
use module level _current_tasks, all_tasks and patch run method.
|
|
"""
|
|
def run(future, *, debug=False):
|
|
loop = asyncio.get_event_loop()
|
|
loop.set_debug(debug)
|
|
return loop.run_until_complete(future)
|
|
|
|
if sys.version_info >= (3, 6, 0):
|
|
asyncio.Task = asyncio.tasks._CTask = asyncio.tasks.Task = \
|
|
asyncio.tasks._PyTask
|
|
asyncio.Future = asyncio.futures._CFuture = asyncio.futures.Future = \
|
|
asyncio.futures._PyFuture
|
|
if sys.version_info < (3, 7, 0):
|
|
asyncio.tasks._current_tasks = asyncio.tasks.Task._current_tasks # noqa
|
|
asyncio.all_tasks = asyncio.tasks.Task.all_tasks # noqa
|
|
if not hasattr(asyncio, '_run_orig'):
|
|
asyncio._run_orig = getattr(asyncio, 'run', None)
|
|
asyncio.run = run
|
|
|
|
|
|
def _patch_loop(loop):
|
|
"""Patch loop to make it reentrent."""
|
|
|
|
def run_forever(self):
|
|
if sys.version_info >= (3, 7, 0):
|
|
set_coro_tracking = self._set_coroutine_origin_tracking
|
|
else:
|
|
set_coro_tracking = self._set_coroutine_wrapper
|
|
|
|
self._check_closed()
|
|
old_thread_id = self._thread_id
|
|
old_running_loop = events._get_running_loop()
|
|
set_coro_tracking(self._debug)
|
|
self._thread_id = threading.get_ident()
|
|
|
|
if self._asyncgens is not None:
|
|
old_agen_hooks = sys.get_asyncgen_hooks()
|
|
sys.set_asyncgen_hooks(
|
|
firstiter=self._asyncgen_firstiter_hook,
|
|
finalizer=self._asyncgen_finalizer_hook)
|
|
try:
|
|
events._set_running_loop(self)
|
|
while True:
|
|
self._run_once()
|
|
if self._stopping:
|
|
break
|
|
finally:
|
|
self._stopping = False
|
|
self._thread_id = old_thread_id
|
|
events._set_running_loop(old_running_loop)
|
|
set_coro_tracking(False)
|
|
if self._asyncgens is not None:
|
|
sys.set_asyncgen_hooks(*old_agen_hooks)
|
|
|
|
def run_until_complete(self, future):
|
|
old_running_loop = events._get_running_loop()
|
|
try:
|
|
self._check_closed()
|
|
events._set_running_loop(self)
|
|
f = asyncio.ensure_future(future, loop=self)
|
|
if f is not future:
|
|
f._log_destroy_pending = False
|
|
while not f.done():
|
|
self._run_once()
|
|
if self._stopping:
|
|
break
|
|
if not f.done():
|
|
raise RuntimeError(
|
|
'Event loop stopped before Future completed.')
|
|
return f.result()
|
|
finally:
|
|
events._set_running_loop(old_running_loop)
|
|
|
|
def _run_once(self):
|
|
"""
|
|
Simplified re-implementation of asyncio's _run_once that
|
|
runs handles as they become ready.
|
|
"""
|
|
now = self.time()
|
|
ready = self._ready
|
|
scheduled = self._scheduled
|
|
while scheduled and scheduled[0]._cancelled:
|
|
heappop(scheduled)
|
|
|
|
timeout = (
|
|
0 if ready or self._stopping
|
|
else min(max(0, scheduled[0]._when - now), 86400) if scheduled
|
|
else 0.01 if self._is_proactorloop
|
|
else None)
|
|
event_list = self._selector.select(timeout)
|
|
self._process_events(event_list)
|
|
|
|
end_time = self.time() + self._clock_resolution
|
|
while scheduled and scheduled[0]._when < end_time:
|
|
handle = heappop(scheduled)
|
|
ready.append(handle)
|
|
|
|
for _ in range(len(ready)):
|
|
if not ready:
|
|
break
|
|
handle = ready.popleft()
|
|
if not handle._cancelled:
|
|
handle._run()
|
|
handle = None
|
|
|
|
def _check_running(self):
|
|
"""Do not throw exception if loop is already running."""
|
|
pass
|
|
|
|
cls = loop.__class__
|
|
cls._run_once_orig = cls._run_once
|
|
cls._run_once = _run_once
|
|
cls._run_forever_orig = cls.run_forever
|
|
cls.run_forever = run_forever
|
|
cls._run_until_complete_orig = cls.run_until_complete
|
|
cls.run_until_complete = run_until_complete
|
|
cls._check_running = _check_running
|
|
cls._check_runnung = _check_running # typo in Python 3.7 source
|
|
cls._nest_patched = True
|
|
cls._is_proactorloop = (
|
|
os.name == 'nt' and issubclass(cls, asyncio.ProactorEventLoop))
|
|
|
|
|
|
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 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
|
|
|
|
|
|
def _patch_handle():
|
|
"""Patch Handle to allow recursive calls."""
|
|
|
|
def update_from_context(ctx):
|
|
"""Copy context ctx to currently active context."""
|
|
for var in ctx:
|
|
var.set(ctx[var])
|
|
|
|
def run(self):
|
|
"""
|
|
Run the callback in a sub-context, then copy any sub-context vars
|
|
over to the Handle's context.
|
|
"""
|
|
try:
|
|
ctx = self._context.copy()
|
|
ctx.run(self._callback, *self._args)
|
|
if ctx:
|
|
self._context.run(update_from_context, ctx)
|
|
except (SystemExit, KeyboardInterrupt):
|
|
raise
|
|
except BaseException as exc:
|
|
cb = format_helpers._format_callback_source(
|
|
self._callback, self._args)
|
|
msg = 'Exception in callback {}'.format(cb)
|
|
context = {
|
|
'message': msg,
|
|
'exception': exc,
|
|
'handle': self,
|
|
}
|
|
if self._source_traceback:
|
|
context['source_traceback'] = self._source_traceback
|
|
self._loop.call_exception_handler(context)
|
|
self = None
|
|
|
|
if sys.version_info >= (3, 7, 0):
|
|
from asyncio import format_helpers
|
|
events.Handle._run = run
|
|
|
|
|
|
def _patch_tornado():
|
|
"""
|
|
If tornado is imported before nest_asyncio, make tornado aware of
|
|
the pure-Python asyncio Future.
|
|
"""
|
|
if 'tornado' in sys.modules:
|
|
import tornado.concurrent as tc
|
|
tc.Future = asyncio.Future
|
|
if asyncio.Future not in tc.FUTURES:
|
|
tc.FUTURES += (asyncio.Future,)
|