| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172 |
- """A trio loop runner."""
- import builtins
- import logging
- import signal
- import threading
- import traceback
- import warnings
- import trio
- class TrioRunner:
- """A trio loop runner."""
- def __init__(self):
- """Initialize the runner."""
- self._cell_cancel_scope = None
- self._trio_token = None
- def initialize(self, kernel, io_loop):
- """Initialize the runner."""
- kernel.shell.set_trio_runner(self)
- kernel.shell.run_line_magic("autoawait", "trio")
- kernel.shell.magics_manager.magics["line"]["autoawait"] = lambda _: warnings.warn(
- "Autoawait isn't allowed in Trio background loop mode.", stacklevel=2
- )
- self._interrupted = False
- bg_thread = threading.Thread(target=io_loop.start, daemon=True, name="TornadoBackground")
- bg_thread.start()
- def interrupt(self, signum, frame):
- """Interrupt the runner."""
- if self._cell_cancel_scope:
- self._cell_cancel_scope.cancel()
- else:
- msg = "Kernel interrupted but no cell is running"
- raise Exception(msg)
- def run(self):
- """Run the loop."""
- old_sig = signal.signal(signal.SIGINT, self.interrupt)
- def log_nursery_exc(exc):
- exc = "\n".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
- logging.error("An exception occurred in a global nursery task.\n%s", exc)
- async def trio_main():
- """Run the main loop."""
- self._trio_token = trio.lowlevel.current_trio_token()
- async with trio.open_nursery() as nursery:
- # TODO This hack prevents the nursery from cancelling all child
- # tasks when an uncaught exception occurs, but it's ugly.
- nursery._add_exc = log_nursery_exc
- builtins.GLOBAL_NURSERY = nursery # type:ignore[attr-defined]
- await trio.sleep_forever()
- trio.run(trio_main)
- signal.signal(signal.SIGINT, old_sig)
- def __call__(self, async_fn):
- """Handle a function call."""
- async def loc(coro):
- """A thread runner context."""
- self._cell_cancel_scope = trio.CancelScope()
- with self._cell_cancel_scope:
- return await coro
- self._cell_cancel_scope = None # type:ignore[unreachable]
- return None
- return trio.from_thread.run(loc, async_fn, trio_token=self._trio_token)
|
