trio_ipython.py

"""
Example of running an embedded IPython shell inside an already-running trio loop with working autoawait (it's handy 
to be able to start an interactive REPL with your application environment fully initialized). This is a full solution 
that works around https://github.com/ipython/ipython/issues/680 (see 
https://gist.github.com/mikenerone/3640fdd450b4ca55ee8df4d4da5a7165 for how simple it *could* be). This bug should be 
fixed IMO (using atexit is a questionable design choice in the first place given the embedding feature of IPython 
IMO). As it is now, the entire IPythonAtExitContext context manager exists only to work around the problem, 
otherwise it would result in an error on process exit when IPython's atexit-registered method calls fail to save the 
input history. 

Note: You may wonder "Why not simply execute and unregister IPython's atexit registrations?" The answer is that they
are bound methods, which can't be unregistered because you can't get a reference to the registered bound method
(referencing the method again gives you a *new* instance of a bound method every time).
"""

from unittest.mock import patch

import IPython
import trio


def trio_embedded_runner(coro):
    return trio.from_thread.run(lambda: coro)


def ipython_worker():
    with IPythonAtExitContext():
        IPython.embed(using=trio_embedded_runner)


async def main():
    print("In trio loop")
    await trio.to_thread.run_sync(ipython_worker)
    print("Exiting trio loop")


class IPythonAtExitContext:

    ipython_modules_with_atexit = [
        "IPython.core.magics.script",
        "IPython.core.application",
        "IPython.core.interactiveshell",
        "IPython.core.history",
        "IPython.utils.io",
    ]

    def __init__(self):
        self._calls = []
        self._patchers = []

    def __enter__(self):
        for module in self.ipython_modules_with_atexit:
            try:
                patcher = patch(module + ".atexit", self)
                patcher.start()
            except (AttributeError, ModuleNotFoundError):
                pass
            else:
                self._patchers.append(patcher)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        for patcher in self._patchers:
            patcher.stop()
        self._patchers.clear()
        cb_exc = None
        for func, args, kwargs in self._calls:
            # noinspection PyBroadException
            try:
                func(*args, **kwargs)
            except Exception as _exc:
                cb_exc = _exc
        self._calls.clear()
        if cb_exc and not exc_type:
            raise cb_exc

    def register(self, func, *args, **kwargs):
        self._calls.append((func, args, kwargs))

    def unregister(self, func):
        self._calls = [call for call in self._calls if call[0] != func]


trio.run(main)