1  """A trio loop runner."""
   2  import builtins
   3  import logging
   4  import signal
   5  import threading
   6  import traceback
   7  import warnings
   8  
   9  import trio
  10  
  11  
  12  class TrioRunner:
  13      """A trio loop runner."""
  14  
  15      def __init__(self):
  16          """Initialize the runner."""
  17          self._cell_cancel_scope = None
  18          self._trio_token = None
  19  
  20      def initialize(self, kernel, io_loop):
  21          """Initialize the runner."""
  22          kernel.shell.set_trio_runner(self)
  23          kernel.shell.run_line_magic("autoawait", "trio")
  24          kernel.shell.magics_manager.magics["line"]["autoawait"] = lambda _: warnings.warn(
  25              "Autoawait isn't allowed in Trio background loop mode.", stacklevel=2
  26          )
  27          self._interrupted = False
  28          bg_thread = threading.Thread(target=io_loop.start, daemon=True, name="TornadoBackground")
  29          bg_thread.start()
  30  
  31      def interrupt(self, signum, frame):
  32          """Interuppt the runner."""
  33          if self._cell_cancel_scope:
  34              self._cell_cancel_scope.cancel()
  35          else:
  36              msg = "Kernel interrupted but no cell is running"
  37              raise Exception(msg)
  38  
  39      def run(self):
  40          """Run the loop."""
  41          old_sig = signal.signal(signal.SIGINT, self.interrupt)
  42  
  43          def log_nursery_exc(exc):
  44              exc = "\n".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
  45              logging.error("An exception occurred in a global nursery task.\n%s", exc)
  46  
  47          async def trio_main():
  48              """Run the main loop."""
  49              self._trio_token = trio.lowlevel.current_trio_token()
  50              async with trio.open_nursery() as nursery:
  51                  # TODO This hack prevents the nursery from cancelling all child
  52                  # tasks when an uncaught exception occurs, but it's ugly.
  53                  nursery._add_exc = log_nursery_exc
  54                  builtins.GLOBAL_NURSERY = nursery  # type:ignore[attr-defined]
  55                  await trio.sleep_forever()
  56  
  57          trio.run(trio_main)
  58          signal.signal(signal.SIGINT, old_sig)
  59  
  60      def __call__(self, async_fn):
  61          """Handle a function call."""
  62  
  63          async def loc(coro):
  64              """A thread runner context."""
  65              self._cell_cancel_scope = trio.CancelScope()
  66              with self._cell_cancel_scope:
  67                  return await coro
  68              self._cell_cancel_scope = None  # type:ignore[unreachable]
  69              return None
  70  
  71          return trio.from_thread.run(loc, async_fn, trio_token=self._trio_token)
  72