"""Default classes for Comm and CommManager, for usage in IPython. """ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. from __future__ import annotations import contextlib import logging import typing as t import uuid from traitlets.utils.importstring import import_item import comm if t.TYPE_CHECKING: from zmq.eventloop.zmqstream import ZMQStream logger = logging.getLogger("Comm") MessageType = t.Dict[str, t.Any] MaybeDict = t.Optional[t.Dict[str, t.Any]] BuffersType = t.Optional[t.List[bytes]] CommCallback = t.Callable[[MessageType], None] CommTargetCallback = t.Callable[["BaseComm", MessageType], None] class BaseComm: """Class for communicating between a Frontend and a Kernel Must be subclassed with a publish_msg method implementation which sends comm messages through the iopub channel. """ def __init__( self, target_name: str = "comm", data: MaybeDict = None, metadata: MaybeDict = None, buffers: BuffersType = None, comm_id: str | None = None, primary: bool = True, target_module: str | None = None, topic: bytes | None = None, _open_data: MaybeDict = None, _close_data: MaybeDict = None, **kwargs: t.Any, ) -> None: super().__init__(**kwargs) self.comm_id = comm_id if comm_id else uuid.uuid4().hex self.primary = primary self.target_name = target_name self.target_module = target_module self.topic = topic if topic else ("comm-%s" % self.comm_id).encode("ascii") self._open_data = _open_data if _open_data else {} self._close_data = _close_data if _close_data else {} self._msg_callback: CommCallback | None = None self._close_callback: CommCallback | None = None self._closed = True 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: str, # noqa: ARG002 data: MaybeDict = None, # noqa: ARG002 metadata: MaybeDict = None, # noqa: ARG002 buffers: BuffersType = None, # noqa: ARG002 **keys: t.Any, # noqa: ARG002 ) -> None: msg = "publish_msg Comm method is not implemented" raise NotImplementedError(msg) def __del__(self) -> None: """trigger close on gc""" with contextlib.suppress(Exception): # any number of things can have gone horribly wrong # when called during interpreter teardown self.close(deleting=True) # publishing messages def open( self, data: MaybeDict = None, metadata: MaybeDict = None, buffers: BuffersType = None ) -> None: """Open the frontend-side version of this comm""" if data is None: data = self._open_data comm_manager = comm.get_comm_manager() if comm_manager is None: msg = "Comms cannot be opened without a comm_manager." # type:ignore[unreachable] raise RuntimeError(msg) 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: MaybeDict = None, metadata: MaybeDict = None, buffers: BuffersType = None, deleting: bool = False, ) -> None: """Close the frontend-side version of this comm""" if self._closed: # only close once return self._closed = True 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 comm.get_comm_manager().unregister_comm(self) def send( self, data: MaybeDict = None, metadata: MaybeDict = None, buffers: BuffersType = None ) -> 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: CommCallback | None) -> None: """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: CommCallback | None) -> None: """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: MessageType) -> None: """Handle a comm_close message""" logger.debug("handle_close[%s](%s)", self.comm_id, msg) if self._close_callback: self._close_callback(msg) def handle_msg(self, msg: MessageType) -> None: """Handle a comm_msg message""" logger.debug("handle_msg[%s](%s)", self.comm_id, msg) if self._msg_callback: from IPython import get_ipython shell = get_ipython() if shell: shell.events.trigger("pre_execute") self._msg_callback(msg) if shell: shell.events.trigger("post_execute") class CommManager: """Default CommManager singleton implementation for Comms in the Kernel""" # Public APIs def __init__(self) -> None: self.comms: dict[str, BaseComm] = {} self.targets: dict[str, CommTargetCallback] = {} def register_target(self, target_name: str, f: CommTargetCallback | str) -> None: """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] = t.cast(CommTargetCallback, f) def unregister_target(self, target_name: str, f: CommTargetCallback) -> CommTargetCallback: # noqa: ARG002 """Unregister a callable registered with register_target""" return self.targets.pop(target_name) def register_comm(self, comm: BaseComm) -> str: """Register a new comm""" comm_id = comm.comm_id self.comms[comm_id] = comm return comm_id def unregister_comm(self, comm: BaseComm) -> None: """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: str) -> BaseComm | None: """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: logger.warning("No such comm: %s", comm_id) if logger.isEnabledFor(logging.DEBUG): # don't create the list of keys if debug messages aren't enabled logger.debug("Current comms: %s", list(self.comms.keys())) return None # Message handlers def comm_open(self, stream: ZMQStream, ident: str, msg: MessageType) -> None: # noqa: ARG002 """Handler for comm_open messages""" from comm import create_comm content = msg["content"] comm_id = content["comm_id"] target_name = content["target_name"] f = self.targets.get(target_name, None) comm = create_comm( comm_id=comm_id, primary=False, target_name=target_name, ) self.register_comm(comm) if f is None: logger.error("No such comm target registered: %s", target_name) else: try: f(comm, msg) return except Exception: logger.error("Exception opening comm with target: %s", target_name, exc_info=True) # Failure. try: comm.close() except Exception: logger.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: ZMQStream, ident: str, msg: MessageType) -> None: # noqa: ARG002 """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: logger.error("Exception in comm_msg for %s", comm_id, exc_info=True) def comm_close(self, stream: ZMQStream, ident: str, msg: MessageType) -> None: # noqa: ARG002 """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: logger.error("Exception in comm_close for %s", comm_id, exc_info=True) __all__ = ["CommManager", "BaseComm"]