Skip to content

Commit

Permalink
feat: support comm package
Browse files Browse the repository at this point in the history
  • Loading branch information
maartenbreddels committed Dec 8, 2022
1 parent 463b4a2 commit cf71d51
Showing 1 changed file with 43 additions and 1 deletion.
44 changes: 43 additions & 1 deletion solara/server/kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import struct
from typing import Set

import ipykernel
import ipykernel.kernelbase
import jupyter_client.session as session
from ipykernel.comm import CommManager
Expand All @@ -17,6 +18,38 @@

logger = logging.getLogger("solara.server.kernel")


ipykernel_version = tuple(map(int, ipykernel.__version__.split(".")))
if ipykernel_version >= (6, 18, 0):
import comm.base_comm

class Comm(comm.base_comm.BaseComm):
def __init__(self, **kwargs) -> None:
self.kernel = ipykernel.kernelbase.Kernel.instance()
super().__init__(**kwargs)

def publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
data = {} if data is None else data
metadata = {} if metadata is None else metadata
content = dict(data=data, comm_id=self.comm_id, **keys)
self.kernel.session.send(
self.kernel.iopub_socket,
msg_type,
content,
metadata=metadata,
parent=self.kernel.get_parent("shell"),
ident=self.topic,
buffers=buffers,
)

comm.create_comm = Comm

def get_comm_manager():
from .app import get_current_context

return get_current_context().kernel.comm_manager

comm.get_comm_manager = get_comm_manager
# from notebook.base.zmqhandlers import serialize_binary_message
# this saves us a depdendency on notebook/jupyter_server when e.g.
# running on pyodide
Expand Down Expand Up @@ -160,7 +193,16 @@ def __init__(self):
# solara/server/kernel.py:111: error: "SessionWebsocket" has no attribute "stream"
# not sure why we cannot reproduce that locally
self.session.stream = self.iopub_socket # type: ignore
self.comm_manager = CommManager(parent=self, kernel=self)
if ipykernel_version >= (6, 18, 0):
# from this version on, ipykernel uses the comm package https://github.com/ipython/ipykernel/pull/973
self.comm_manager = CommManager(parent=self, kernel=self)
import ipywidgets.widgets.widget

if hasattr(ipywidgets.widgets.widget, "Comm"):
ipywidgets.widgets.widget.Comm = Comm
ipywidgets.widgets.widget.Widget.comm.klass = Comm
else:
self.comm_manager = CommManager(parent=self, kernel=self)
self.shell = None
self.log = logging.getLogger("fake")

Expand Down

0 comments on commit cf71d51

Please sign in to comment.