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