274 lines
8.7 KiB
Python
274 lines
8.7 KiB
Python
|
"""
|
||
|
|
Based on
|
||
|
|
https://github.com/jupyter/notebook/blob/master/notebook/static/services/kernels/comm.js
|
||
|
|
https://github.com/ipython/ipykernel/blob/master/ipykernel/comm/manager.py
|
||
|
|
https://github.com/ipython/ipykernel/blob/master/ipykernel/comm/comm.py
|
||
|
|
|
||
|
|
|
||
|
|
Which are distributed under the terms of the Modified BSD License.
|
||
|
|
"""
|
||
|
|
import logging
|
||
|
|
|
||
|
|
from traitlets.config import LoggingConfigurable
|
||
|
|
|
||
|
|
from ipython_genutils.importstring import import_item
|
||
|
|
|
||
|
|
import uuid
|
||
|
|
|
||
|
|
from qtpy import QtCore
|
||
|
|
from qtconsole.util import MetaQObjectHasTraits, SuperQObject
|
||
|
|
|
||
|
|
|
||
|
|
class CommManager(MetaQObjectHasTraits(
|
||
|
|
'NewBase', (LoggingConfigurable, SuperQObject), {})):
|
||
|
|
"""
|
||
|
|
Manager for Comms in the Frontend
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, kernel_client, *args, **kwargs):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
self.comms = {}
|
||
|
|
self.targets = {}
|
||
|
|
if kernel_client:
|
||
|
|
self.init_kernel_client(kernel_client)
|
||
|
|
|
||
|
|
def init_kernel_client(self, kernel_client):
|
||
|
|
"""
|
||
|
|
connect the kernel, and register message handlers
|
||
|
|
"""
|
||
|
|
self.kernel_client = kernel_client
|
||
|
|
kernel_client.iopub_channel.message_received.connect(self._dispatch)
|
||
|
|
|
||
|
|
@QtCore.Slot(object)
|
||
|
|
def _dispatch(self, msg):
|
||
|
|
"""Dispatch messages"""
|
||
|
|
msg_type = msg['header']['msg_type']
|
||
|
|
handled_msg_types = ['comm_open', 'comm_msg', 'comm_close']
|
||
|
|
if msg_type in handled_msg_types:
|
||
|
|
getattr(self, msg_type)(msg)
|
||
|
|
|
||
|
|
def new_comm(self, target_name, data=None, metadata=None,
|
||
|
|
comm_id=None, buffers=None):
|
||
|
|
"""
|
||
|
|
Create a new Comm, register it, and open its Kernel-side counterpart
|
||
|
|
Mimics the auto-registration in `Comm.__init__` in the Jupyter Comm.
|
||
|
|
|
||
|
|
argument comm_id is optional
|
||
|
|
"""
|
||
|
|
comm = Comm(target_name, self.kernel_client, comm_id)
|
||
|
|
self.register_comm(comm)
|
||
|
|
try:
|
||
|
|
comm.open(data, metadata, buffers)
|
||
|
|
except Exception:
|
||
|
|
self.unregister_comm(comm)
|
||
|
|
raise
|
||
|
|
return comm
|
||
|
|
|
||
|
|
def register_target(self, target_name, f):
|
||
|
|
"""Register a callable f for a given target name
|
||
|
|
|
||
|
|
f will be called with two arguments when a comm_open message is
|
||
|
|
received with `target`:
|
||
|
|
|
||
|
|
- the Comm instance
|
||
|
|
- the `comm_open` message itself.
|
||
|
|
|
||
|
|
f can be a Python callable or an import string for one.
|
||
|
|
"""
|
||
|
|
if isinstance(f, str):
|
||
|
|
f = import_item(f)
|
||
|
|
|
||
|
|
self.targets[target_name] = f
|
||
|
|
|
||
|
|
def unregister_target(self, target_name, f):
|
||
|
|
"""Unregister a callable registered with register_target"""
|
||
|
|
return self.targets.pop(target_name)
|
||
|
|
|
||
|
|
def register_comm(self, comm):
|
||
|
|
"""Register a new comm"""
|
||
|
|
comm_id = comm.comm_id
|
||
|
|
comm.kernel_client = self.kernel_client
|
||
|
|
self.comms[comm_id] = comm
|
||
|
|
comm.sig_is_closing.connect(self.unregister_comm)
|
||
|
|
return comm_id
|
||
|
|
|
||
|
|
@QtCore.Slot(object)
|
||
|
|
def unregister_comm(self, comm):
|
||
|
|
"""Unregister a comm, and close its counterpart."""
|
||
|
|
# unlike get_comm, this should raise a KeyError
|
||
|
|
comm.sig_is_closing.disconnect(self.unregister_comm)
|
||
|
|
self.comms.pop(comm.comm_id)
|
||
|
|
|
||
|
|
def get_comm(self, comm_id, closing=False):
|
||
|
|
"""Get a comm with a particular id
|
||
|
|
|
||
|
|
Returns the comm if found, otherwise None.
|
||
|
|
|
||
|
|
This will not raise an error,
|
||
|
|
it will log messages if the comm cannot be found.
|
||
|
|
If the comm is closing, it might already have closed,
|
||
|
|
so this is ignored.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
return self.comms[comm_id]
|
||
|
|
except KeyError:
|
||
|
|
if closing:
|
||
|
|
return
|
||
|
|
self.log.warning("No such comm: %s", comm_id)
|
||
|
|
# don't create the list of keys if debug messages aren't enabled
|
||
|
|
if self.log.isEnabledFor(logging.DEBUG):
|
||
|
|
self.log.debug("Current comms: %s", list(self.comms.keys()))
|
||
|
|
|
||
|
|
# comm message handlers
|
||
|
|
def comm_open(self, msg):
|
||
|
|
"""Handler for comm_open messages"""
|
||
|
|
content = msg['content']
|
||
|
|
comm_id = content['comm_id']
|
||
|
|
target_name = content['target_name']
|
||
|
|
f = self.targets.get(target_name, None)
|
||
|
|
|
||
|
|
comm = Comm(target_name, self.kernel_client, comm_id)
|
||
|
|
self.register_comm(comm)
|
||
|
|
|
||
|
|
if f is None:
|
||
|
|
self.log.error("No such comm target registered: %s", target_name)
|
||
|
|
else:
|
||
|
|
try:
|
||
|
|
f(comm, msg)
|
||
|
|
return
|
||
|
|
except Exception:
|
||
|
|
self.log.error("Exception opening comm with target: %s",
|
||
|
|
target_name, exc_info=True)
|
||
|
|
|
||
|
|
# Failure.
|
||
|
|
try:
|
||
|
|
comm.close()
|
||
|
|
except Exception:
|
||
|
|
self.log.error(
|
||
|
|
"Could not close comm during `comm_open` failure "
|
||
|
|
"clean-up. The comm may not have been opened yet.""",
|
||
|
|
exc_info=True)
|
||
|
|
|
||
|
|
def comm_close(self, msg):
|
||
|
|
"""Handler for comm_close messages"""
|
||
|
|
content = msg['content']
|
||
|
|
comm_id = content['comm_id']
|
||
|
|
comm = self.get_comm(comm_id, closing=True)
|
||
|
|
if comm is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
self.unregister_comm(comm)
|
||
|
|
|
||
|
|
try:
|
||
|
|
comm.handle_close(msg)
|
||
|
|
except Exception:
|
||
|
|
self.log.error('Exception in comm_close for %s', comm_id,
|
||
|
|
exc_info=True)
|
||
|
|
|
||
|
|
def comm_msg(self, msg):
|
||
|
|
"""Handler for comm_msg messages"""
|
||
|
|
content = msg['content']
|
||
|
|
comm_id = content['comm_id']
|
||
|
|
comm = self.get_comm(comm_id)
|
||
|
|
if comm is None:
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
comm.handle_msg(msg)
|
||
|
|
except Exception:
|
||
|
|
self.log.error('Exception in comm_msg for %s', comm_id,
|
||
|
|
exc_info=True)
|
||
|
|
|
||
|
|
|
||
|
|
class Comm(MetaQObjectHasTraits(
|
||
|
|
'NewBase', (LoggingConfigurable, SuperQObject), {})):
|
||
|
|
"""
|
||
|
|
Comm base class
|
||
|
|
"""
|
||
|
|
sig_is_closing = QtCore.Signal(object)
|
||
|
|
|
||
|
|
def __init__(self, target_name, kernel_client, comm_id=None,
|
||
|
|
msg_callback=None, close_callback=None):
|
||
|
|
"""
|
||
|
|
Create a new comm. Must call open to use.
|
||
|
|
"""
|
||
|
|
super().__init__(target_name=target_name)
|
||
|
|
self.target_name = target_name
|
||
|
|
self.kernel_client = kernel_client
|
||
|
|
if comm_id is None:
|
||
|
|
comm_id = uuid.uuid1().hex
|
||
|
|
self.comm_id = comm_id
|
||
|
|
self._msg_callback = msg_callback
|
||
|
|
self._close_callback = close_callback
|
||
|
|
self._send_channel = self.kernel_client.shell_channel
|
||
|
|
|
||
|
|
def _send_msg(self, msg_type, content, data, metadata, buffers):
|
||
|
|
"""
|
||
|
|
Send a message on the shell channel.
|
||
|
|
"""
|
||
|
|
if data is None:
|
||
|
|
data = {}
|
||
|
|
if content is None:
|
||
|
|
content = {}
|
||
|
|
content['comm_id'] = self.comm_id
|
||
|
|
content['data'] = data
|
||
|
|
|
||
|
|
msg = self.kernel_client.session.msg(
|
||
|
|
msg_type, content, metadata=metadata)
|
||
|
|
if buffers:
|
||
|
|
msg['buffers'] = buffers
|
||
|
|
return self._send_channel.send(msg)
|
||
|
|
|
||
|
|
# methods for sending messages
|
||
|
|
def open(self, data=None, metadata=None, buffers=None):
|
||
|
|
"""Open the kernel-side version of this comm"""
|
||
|
|
return self._send_msg(
|
||
|
|
'comm_open', {'target_name': self.target_name},
|
||
|
|
data, metadata, buffers)
|
||
|
|
|
||
|
|
def send(self, data=None, metadata=None, buffers=None):
|
||
|
|
"""Send a message to the kernel-side version of this comm"""
|
||
|
|
return self._send_msg(
|
||
|
|
'comm_msg', {}, data, metadata, buffers)
|
||
|
|
|
||
|
|
def close(self, data=None, metadata=None, buffers=None):
|
||
|
|
"""Close the kernel-side version of this comm"""
|
||
|
|
self.sig_is_closing.emit(self)
|
||
|
|
return self._send_msg(
|
||
|
|
'comm_close', {}, data, metadata, buffers)
|
||
|
|
|
||
|
|
# methods for registering callbacks for incoming messages
|
||
|
|
|
||
|
|
def on_msg(self, callback):
|
||
|
|
"""Register a callback for comm_msg
|
||
|
|
|
||
|
|
Will be called with the `data` of any comm_msg messages.
|
||
|
|
|
||
|
|
Call `on_msg(None)` to disable an existing callback.
|
||
|
|
"""
|
||
|
|
self._msg_callback = callback
|
||
|
|
|
||
|
|
def on_close(self, callback):
|
||
|
|
"""Register a callback for comm_close
|
||
|
|
|
||
|
|
Will be called with the `data` of the close message.
|
||
|
|
|
||
|
|
Call `on_close(None)` to disable an existing callback.
|
||
|
|
"""
|
||
|
|
self._close_callback = callback
|
||
|
|
|
||
|
|
# methods for handling incoming messages
|
||
|
|
def handle_msg(self, msg):
|
||
|
|
"""Handle a comm_msg message"""
|
||
|
|
self.log.debug("handle_msg[%s](%s)", self.comm_id, msg)
|
||
|
|
if self._msg_callback:
|
||
|
|
return self._msg_callback(msg)
|
||
|
|
|
||
|
|
def handle_close(self, msg):
|
||
|
|
"""Handle a comm_close message"""
|
||
|
|
self.log.debug("handle_close[%s](%s)", self.comm_id, msg)
|
||
|
|
if self._close_callback:
|
||
|
|
return self._close_callback(msg)
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = ['CommManager']
|