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