Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Extract the Comm Python package #973

Merged
merged 2 commits into from
Nov 18, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions ipykernel/comm/__init__.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,5 @@
from .comm import * # noqa
from .manager import * # noqa
__all__ = ["Comm", "CommManager"]

from comm.base_comm import CommManager # noqa

from .comm import Comm # noqa
167 changes: 13 additions & 154 deletions ipykernel/comm/comm.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,71 +3,32 @@
# Copyright (c) IPython Development Team.
# Distributed under the terms of the Modified BSD License.

import uuid

from traitlets import Any, Bool, Bytes, Dict, Instance, Unicode, default
from traitlets.config import LoggingConfigurable
from comm.base_comm import BaseComm

from ipykernel.jsonutil import json_clean
from ipykernel.kernelbase import Kernel


class Comm(LoggingConfigurable):
class Comm(BaseComm):
"""Class for communicating between a Frontend and a Kernel"""

kernel = Instance("ipykernel.kernelbase.Kernel", allow_none=True)

@default("kernel")
def _default_kernel(self):
if Kernel.initialized():
return Kernel.instance()

comm_id = Unicode()

@default("comm_id")
def _default_comm_id(self):
return uuid.uuid4().hex

primary = Bool(True, help="Am I the primary or secondary Comm?")

target_name = Unicode("comm")
target_module = Unicode(
None,
allow_none=True,
help="""requirejs module from
which to load comm target.""",
)

topic = Bytes()

@default("topic")
def _default_topic(self):
return ("comm-%s" % self.comm_id).encode("ascii")

_open_data = Dict(help="data dict, if any, to be included in comm_open")
_close_data = Dict(help="data dict, if any, to be included in comm_close")

_msg_callback = Any()
_close_callback = Any()
def __init__(self, *args, **kwargs):
self.kernel = None

_closed = Bool(True)
super().__init__(*args, **kwargs)

def __init__(self, target_name="", data=None, metadata=None, buffers=None, **kwargs):
if target_name:
kwargs["target_name"] = target_name
super().__init__(**kwargs)
if self.kernel:
if self.primary:
# I am primary, open my peer.
self.open(data=data, metadata=metadata, buffers=buffers)
else:
self._closed = False

def _publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
def publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
"""Helper for sending a comm message on IOPub"""
if not Kernel.initialized():
return

data = {} if data is None else data
metadata = {} if metadata is None else metadata
content = json_clean(dict(data=data, comm_id=self.comm_id, **keys))

if self.kernel is None:
self.kernel = Kernel.instance()

self.kernel.session.send(
self.kernel.iopub_socket,
msg_type,
Expand All @@ -78,107 +39,5 @@ def _publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys)
buffers=buffers,
)

def __del__(self):
"""trigger close on gc"""
self.close(deleting=True)

# publishing messages

def open(self, data=None, metadata=None, buffers=None):
"""Open the frontend-side version of this comm"""
if data is None:
data = self._open_data
comm_manager = getattr(self.kernel, "comm_manager", None)
if comm_manager is None:
raise RuntimeError(
"Comms cannot be opened without a kernel "
"and a comm_manager attached to that kernel."
)

comm_manager.register_comm(self)
try:
self._publish_msg(
"comm_open",
data=data,
metadata=metadata,
buffers=buffers,
target_name=self.target_name,
target_module=self.target_module,
)
self._closed = False
except Exception:
comm_manager.unregister_comm(self)
raise

def close(self, data=None, metadata=None, buffers=None, deleting=False):
"""Close the frontend-side version of this comm"""
if self._closed:
# only close once
return
self._closed = True
# nothing to send if we have no kernel
# can be None during interpreter cleanup
if not self.kernel:
return
if data is None:
data = self._close_data
self._publish_msg(
"comm_close",
data=data,
metadata=metadata,
buffers=buffers,
)
if not deleting:
# If deleting, the comm can't be registered
self.kernel.comm_manager.unregister_comm(self)

def send(self, data=None, metadata=None, buffers=None):
"""Send a message to the frontend-side version of this comm"""
self._publish_msg(
"comm_msg",
data=data,
metadata=metadata,
buffers=buffers,
)

# registering callbacks

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

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

# handling of incoming messages

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:
self._close_callback(msg)

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:
shell = self.kernel.shell
if shell:
shell.events.trigger("pre_execute")
self._msg_callback(msg)
if shell:
shell.events.trigger("post_execute")


__all__ = ["Comm"]
129 changes: 1 addition & 128 deletions ipykernel/comm/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,131 +3,4 @@
# Copyright (c) IPython Development Team.
# Distributed under the terms of the Modified BSD License.

import logging

from traitlets import Dict, Instance
from traitlets.config import LoggingConfigurable
from traitlets.utils.importstring import import_item

from .comm import Comm


class CommManager(LoggingConfigurable):
"""Manager for Comms in the Kernel"""

kernel = Instance("ipykernel.kernelbase.Kernel")
comms = Dict()
targets = Dict()

# Public APIs

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 = self.kernel
self.comms[comm_id] = comm
return comm_id

def unregister_comm(self, comm):
"""Unregister a comm, and close its counterpart"""
# unlike get_comm, this should raise a KeyError
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)
if self.log.isEnabledFor(logging.DEBUG):
# don't create the list of keys if debug messages aren't enabled
self.log.debug("Current comms: %s", list(self.comms.keys()))

# Message handlers
def comm_open(self, stream, ident, 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(
comm_id=comm_id,
primary=False,
target_name=target_name,
)
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_msg(self, stream, ident, 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)

def comm_close(self, stream, ident, 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.comms[comm_id]._closed = True
del self.comms[comm_id]

try:
comm.handle_close(msg)
except Exception:
self.log.error("Exception in comm_close for %s", comm_id, exc_info=True)


__all__ = ["CommManager"]
from comm.base_comm import CommManager # noqa
13 changes: 11 additions & 2 deletions ipykernel/ipkernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,13 @@
from contextlib import contextmanager
from functools import partial

import comm
from IPython.core import release
from IPython.utils.tokenutil import line_at_cursor, token_at_cursor
from traitlets import Any, Bool, Instance, List, Type, observe, observe_compat
from zmq.eventloop.zmqstream import ZMQStream

from .comm import CommManager
from .comm import Comm
from .compiler import XCachingCompiler
from .debugger import Debugger, _is_debugpy_available
from .eventloops import _use_appnope
Expand All @@ -39,6 +40,14 @@
_EXPERIMENTAL_KEY_NAME = "_jupyter_types_experimental"


def create_comm(*args, **kwargs):
"""Create a new Comm."""
return Comm(*args, **kwargs)


comm.create_comm = create_comm


class IPythonKernel(KernelBase):
shell = Instance("IPython.core.interactiveshell.InteractiveShellABC", allow_none=True)
shell_class = Type(ZMQInteractiveShell)
Expand Down Expand Up @@ -101,7 +110,7 @@ def __init__(self, **kwargs):
self.shell.display_pub.session = self.session
self.shell.display_pub.pub_socket = self.iopub_socket

self.comm_manager = CommManager(parent=self, kernel=self)
self.comm_manager = comm.get_comm_manager()

self.shell.configurables.append(self.comm_manager)
comm_msg_types = ["comm_open", "comm_msg", "comm_close"]
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ requires-python = ">=3.8"
dependencies = [
"debugpy>=1.0",
"ipython>=7.23.1",
"comm>=0.1",
"traitlets>=5.1.0",
"jupyter_client>=6.1.12",
"tornado>=6.1",
Expand Down