allow for loading / unloading as an ipython extension

core/ipyflow/__init__.py

@@ -1,6 +1,9 @@
 # -*- coding: utf-8 -*-
+from IPython import InteractiveShell
+
 import ipyflow.api
 from ipyflow.api import *
+from ipyflow.kernel import IPyflowKernel
 
 
 # Jupyter Extension points
@@ -22,6 +25,20 @@ def load_jupyter_server_extension(nbapp):
     pass
 
 
+def load_ipython_extension(ipy: InteractiveShell) -> None:
+    cur_kernel_cls = ipy.kernel.__class__  # type: ignore
+    if cur_kernel_cls is IPyflowKernel:
+        IPyflowKernel.replacement_class = None  # type: ignore
+        return
+    IPyflowKernel.inject(prev_kernel_class=cur_kernel_cls)  # type: ignore
+
+
+def unload_ipython_extension(ipy: InteractiveShell) -> None:
+    assert isinstance(ipy.kernel, IPyflowKernel)  # type: ignore
+    assert IPyflowKernel.prev_kernel_class is not None  # type: ignore
+    IPyflowKernel.replacement_class = IPyflowKernel.prev_kernel_class  # type: ignore
+
+
 from . import _version
 __version__ = _version.get_versions()['version']
 

core/ipyflow/kernel/kernel.py

@@ -90,11 +90,14 @@ def should_patch_meta_path(self) -> bool:
 
 class PyccoloKernelMixin(PyccoloKernelHooks):
     def __init__(self, **kwargs) -> None:
-        self.settings: PyccoloKernelSettings = PyccoloKernelSettings(
-            store_history=kwargs.pop("store_history", True)
-        )
+        store_history = kwargs.pop("store_history", True)
         super().__init__(**kwargs)
+        self._initialize(store_history=store_history)
 
+    def _initialize(self, store_history: bool = True) -> None:
+        self.settings: PyccoloKernelSettings = PyccoloKernelSettings(
+            store_history=store_history
+        )
         self.tee_output_tracer = OutputRecorder.instance()
         self.registered_tracers: List[Type[pyc.BaseTracer]] = [
             OutputRecorder,
@@ -363,13 +366,43 @@ def make_zmq_kernel_class(cls, name: str) -> Type[IPythonKernel]:
         class ZMQKernel(cls, IPythonKernel):  # type: ignore
             implementation = "kernel"
             implementation_version = __version__
+            prev_kernel_class: Optional[Type[IPythonKernel]] = None
+            replacement_class: Optional[Type[IPythonKernel]] = None
 
-            def __init__(self, **kwargs):
+            def __init__(self, **kwargs) -> None:
                 super().__init__(**kwargs)
+                # this needs to happen after IPythonKernel.__init__ completes,
+                # which means that we cannot put it in _initialize(), since at
+                # that point only PyccoloKernelHooks.__init__ will be finished
+                self.after_init_class()
+
+            def _initialize(self, **kwargs) -> None:
+                super()._initialize(**kwargs)
                 from ipyflow.kernel import patched_nest_asyncio
 
                 patched_nest_asyncio.apply()
-                self.after_init_class()
+
+            @classmethod
+            def inject(
+                zmq_kernel_class, prev_kernel_class: Type[IPythonKernel]
+            ) -> None:
+                ipy = get_ipython()
+                kernel = ipy.kernel
+                kernel.__class__ = zmq_kernel_class
+                if zmq_kernel_class.prev_kernel_class is None:
+                    kernel._initialize()
+                    kernel.after_init_class()
+                    for subclass in singletons.IPyflowKernel._walk_mro():
+                        subclass._instance = kernel
+                zmq_kernel_class.prev_kernel_class = prev_kernel_class
+                cells()._cell_counter = ipy.execution_count
+
+            @classmethod
+            def _maybe_eject(zmq_kernel_class) -> None:
+                if zmq_kernel_class.replacement_class is None:
+                    return
+                get_ipython().kernel.__class__ = zmq_kernel_class.replacement_class
+                zmq_kernel_class.replacement_class = None
 
             def init_metadata(self, parent):
                 """
@@ -409,6 +442,7 @@ async def _run_cell_func(cell):
                         ret = await self.pyc_execute(code, True, _run_cell_func)
                     if ret["status"] == "error":
                         self.on_exception(ret["ename"])
+                    self._maybe_eject()
                     return ret
 
             else:
@@ -439,6 +473,7 @@ async def _run_cell_func(cell):
                     )
                     if ret["status"] == "error":
                         self.on_exception(ret["ename"])
+                    self._maybe_eject()
                     return ret
 
         ZMQKernel.__name__ = name