""" 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 from ipython_genutils.py3compat import string_types import uuid from qtconsole.qt 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(CommManager, self).__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, string_types): 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): """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. """ try: return self.comms[comm_id] except KeyError: 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) 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(Comm, self).__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']