1  """Patch asyncio to allow nested event loops."""
   2  
   3  import asyncio
   4  import asyncio.events as events
   5  import os
   6  import sys
   7  import threading
   8  from contextlib import contextmanager, suppress
   9  from heapq import heappop
  10  
  11  
  12  def apply(loop=None):
  13      """Patch asyncio to make its event loop reentrant."""
  14      _patch_asyncio()
  15      _patch_policy()
  16      _patch_tornado()
  17  
  18      loop = loop or asyncio.get_event_loop()
  19      _patch_loop(loop)
  20  
  21  
  22  def _patch_asyncio():
  23      """Patch asyncio module to use pure Python tasks and futures."""
  24  
  25      def run(main, *, debug=False):
  26          loop = asyncio.get_event_loop()
  27          loop.set_debug(debug)
  28          task = asyncio.ensure_future(main)
  29          try:
  30              return loop.run_until_complete(task)
  31          finally:
  32              if not task.done():
  33                  task.cancel()
  34                  with suppress(asyncio.CancelledError):
  35                      loop.run_until_complete(task)
  36  
  37      def _get_event_loop(stacklevel=3):
  38          loop = events._get_running_loop()
  39          if loop is None:
  40              loop = events.get_event_loop_policy().get_event_loop()
  41          return loop
  42  
  43      # Use module level _current_tasks, all_tasks and patch run method.
  44      if hasattr(asyncio, '_nest_patched'):
  45          return
  46      if sys.version_info >= (3, 6, 0):
  47          asyncio.Task = asyncio.tasks._CTask = asyncio.tasks.Task = \
  48              asyncio.tasks._PyTask
  49          asyncio.Future = asyncio.futures._CFuture = asyncio.futures.Future = \
  50              asyncio.futures._PyFuture
  51      if sys.version_info < (3, 7, 0):
  52          asyncio.tasks._current_tasks = asyncio.tasks.Task._current_tasks
  53          asyncio.all_tasks = asyncio.tasks.Task.all_tasks
  54      if sys.version_info >= (3, 9, 0):
  55          events._get_event_loop = events.get_event_loop = \
  56              asyncio.get_event_loop = _get_event_loop
  57      asyncio.run = run
  58      asyncio._nest_patched = True
  59  
  60  
  61  def _patch_policy():
  62      """Patch the policy to always return a patched loop."""
  63  
  64      def get_event_loop(self):
  65          if self._local._loop is None:
  66              loop = self.new_event_loop()
  67              _patch_loop(loop)
  68              self.set_event_loop(loop)
  69          return self._local._loop
  70  
  71      policy = events.get_event_loop_policy()
  72      policy.__class__.get_event_loop = get_event_loop
  73  
  74  
  75  def _patch_loop(loop):
  76      """Patch loop to make it reentrant."""
  77  
  78      def run_forever(self):
  79          with manage_run(self), manage_asyncgens(self):
  80              while True:
  81                  self._run_once()
  82                  if self._stopping:
  83                      break
  84          self._stopping = False
  85  
  86      def run_until_complete(self, future):
  87          with manage_run(self):
  88              f = asyncio.ensure_future(future, loop=self)
  89              if f is not future:
  90                  f._log_destroy_pending = False
  91              while not f.done():
  92                  self._run_once()
  93                  if self._stopping:
  94                      break
  95              if not f.done():
  96                  raise RuntimeError(
  97                      'Event loop stopped before Future completed.')
  98              return f.result()
  99  
 100      def _run_once(self):
 101          """
 102          Simplified re-implementation of asyncio's _run_once that
 103          runs handles as they become ready.
 104          """
 105          ready = self._ready
 106          scheduled = self._scheduled
 107          while scheduled and scheduled[0]._cancelled:
 108              heappop(scheduled)
 109  
 110          timeout = (
 111              0 if ready or self._stopping
 112              else min(max(
 113                  scheduled[0]._when - self.time(), 0), 86400) if scheduled
 114              else None)
 115          event_list = self._selector.select(timeout)
 116          self._process_events(event_list)
 117  
 118          end_time = self.time() + self._clock_resolution
 119          while scheduled and scheduled[0]._when < end_time:
 120              handle = heappop(scheduled)
 121              ready.append(handle)
 122  
 123          for _ in range(len(ready)):
 124              if not ready:
 125                  break
 126              handle = ready.popleft()
 127              if not handle._cancelled:
 128                  # preempt the current task so that that checks in
 129                  # Task.__step do not raise
 130                  curr_task = curr_tasks.pop(self, None)
 131  
 132                  try:
 133                      handle._run()
 134                  finally:
 135                      # restore the current task
 136                      if curr_task is not None:
 137                          curr_tasks[self] = curr_task
 138  
 139          handle = None
 140  
 141      @contextmanager
 142      def manage_run(self):
 143          """Set up the loop for running."""
 144          self._check_closed()
 145          old_thread_id = self._thread_id
 146          old_running_loop = events._get_running_loop()
 147          try:
 148              self._thread_id = threading.get_ident()
 149              events._set_running_loop(self)
 150              self._num_runs_pending += 1
 151              if self._is_proactorloop:
 152                  if self._self_reading_future is None:
 153                      self.call_soon(self._loop_self_reading)
 154              yield
 155          finally:
 156              self._thread_id = old_thread_id
 157              events._set_running_loop(old_running_loop)
 158              self._num_runs_pending -= 1
 159              if self._is_proactorloop:
 160                  if (self._num_runs_pending == 0
 161                          and self._self_reading_future is not None):
 162                      ov = self._self_reading_future._ov
 163                      self._self_reading_future.cancel()
 164                      if ov is not None:
 165                          self._proactor._unregister(ov)
 166                      self._self_reading_future = None
 167  
 168      @contextmanager
 169      def manage_asyncgens(self):
 170          if not hasattr(sys, 'get_asyncgen_hooks'):
 171              # Python version is too old.
 172              return
 173          old_agen_hooks = sys.get_asyncgen_hooks()
 174          try:
 175              self._set_coroutine_origin_tracking(self._debug)
 176              if self._asyncgens is not None:
 177                  sys.set_asyncgen_hooks(
 178                      firstiter=self._asyncgen_firstiter_hook,
 179                      finalizer=self._asyncgen_finalizer_hook)
 180              yield
 181          finally:
 182              self._set_coroutine_origin_tracking(False)
 183              if self._asyncgens is not None:
 184                  sys.set_asyncgen_hooks(*old_agen_hooks)
 185  
 186      def _check_running(self):
 187          """Do not throw exception if loop is already running."""
 188          pass
 189  
 190      if hasattr(loop, '_nest_patched'):
 191          return
 192      if not isinstance(loop, asyncio.BaseEventLoop):
 193          raise ValueError('Can\'t patch loop of type %s' % type(loop))
 194      cls = loop.__class__
 195      cls.run_forever = run_forever
 196      cls.run_until_complete = run_until_complete
 197      cls._run_once = _run_once
 198      cls._check_running = _check_running
 199      cls._check_runnung = _check_running  # typo in Python 3.7 source
 200      cls._num_runs_pending = 1 if loop.is_running() else 0
 201      cls._is_proactorloop = (
 202          os.name == 'nt' and issubclass(cls, asyncio.ProactorEventLoop))
 203      if sys.version_info < (3, 7, 0):
 204          cls._set_coroutine_origin_tracking = cls._set_coroutine_wrapper
 205      curr_tasks = asyncio.tasks._current_tasks \
 206          if sys.version_info >= (3, 7, 0) else asyncio.Task._current_tasks
 207      cls._nest_patched = True
 208  
 209  
 210  def _patch_tornado():
 211      """
 212      If tornado is imported before nest_asyncio, make tornado aware of
 213      the pure-Python asyncio Future.
 214      """
 215      if 'tornado' in sys.modules:
 216          import tornado.concurrent as tc  # type: ignore
 217          tc.Future = asyncio.Future
 218          if asyncio.Future not in tc.FUTURES:
 219              tc.FUTURES += (asyncio.Future,)
 220