manager.py
150 lines
| 4.7 KiB
| text/x-python
|
PythonLexer
|
r13195 | """Base class to manage comms""" | ||
|
r17099 | # Copyright (c) IPython Development Team. | ||
# Distributed under the terms of the Modified BSD License. | ||||
|
r13195 | |||
|
r13202 | import sys | ||
|
r13195 | from IPython.config import LoggingConfigurable | ||
from IPython.core.prompts import LazyEvaluate | ||||
from IPython.core.getipython import get_ipython | ||||
from IPython.utils.importstring import import_item | ||||
|
r13390 | from IPython.utils.py3compat import string_types | ||
|
r13195 | from IPython.utils.traitlets import Instance, Unicode, Dict, Any | ||
from .comm import Comm | ||||
def lazy_keys(dikt): | ||||
"""Return lazy-evaluated string representation of a dictionary's keys | ||||
Key list is only constructed if it will actually be used. | ||||
Used for debug-logging. | ||||
""" | ||||
return LazyEvaluate(lambda d: list(d.keys())) | ||||
class CommManager(LoggingConfigurable): | ||||
"""Manager for Comms in the Kernel""" | ||||
shell = Instance('IPython.core.interactiveshell.InteractiveShellABC') | ||||
def _shell_default(self): | ||||
return get_ipython() | ||||
iopub_socket = Any() | ||||
def _iopub_socket_default(self): | ||||
|
r13199 | return self.shell.kernel.iopub_socket | ||
|
r13195 | session = Instance('IPython.kernel.zmq.session.Session') | ||
def _session_default(self): | ||||
if self.shell is None: | ||||
return | ||||
|
r13199 | return self.shell.kernel.session | ||
|
r13195 | |||
comms = Dict() | ||||
targets = Dict() | ||||
# Public APIs | ||||
|
r13204 | def register_target(self, target_name, f): | ||
"""Register a callable f for a given target name | ||||
|
r13195 | |||
|
r13227 | f will be called with two arguments when a comm_open message is received with `target`: | ||
- the Comm instance | ||||
- the `comm_open` message itself. | ||||
|
r13195 | |||
f can be a Python callable or an import string for one. | ||||
""" | ||||
|
r13390 | if isinstance(f, string_types): | ||
|
r13195 | f = import_item(f) | ||
|
r13204 | self.targets[target_name] = f | ||
|
r13195 | |||
|
r13226 | def unregister_target(self, target_name, f): | ||
"""Unregister a callable registered with register_target""" | ||||
return self.targets.pop(target_name); | ||||
|
r13227 | |||
|
r13195 | def register_comm(self, comm): | ||
"""Register a new comm""" | ||||
comm_id = comm.comm_id | ||||
comm.shell = self.shell | ||||
comm.iopub_socket = self.iopub_socket | ||||
self.comms[comm_id] = comm | ||||
return comm_id | ||||
def unregister_comm(self, comm_id): | ||||
"""Unregister a comm, and close its counterpart""" | ||||
# unlike get_comm, this should raise a KeyError | ||||
comm = self.comms.pop(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. | ||||
""" | ||||
if comm_id not in self.comms: | ||||
self.log.error("No such comm: %s", comm_id) | ||||
self.log.debug("Current comms: %s", lazy_keys(self.comms)) | ||||
return | ||||
# call, because we store weakrefs | ||||
comm = self.comms[comm_id] | ||||
return comm | ||||
# Message handlers | ||||
def comm_open(self, stream, ident, msg): | ||||
"""Handler for comm_open messages""" | ||||
content = msg['content'] | ||||
comm_id = content['comm_id'] | ||||
|
r13204 | target_name = content['target_name'] | ||
f = self.targets.get(target_name, None) | ||||
|
r13195 | comm = Comm(comm_id=comm_id, | ||
shell=self.shell, | ||||
iopub_socket=self.iopub_socket, | ||||
primary=False, | ||||
) | ||||
|
r13204 | if f is None: | ||
self.log.error("No such comm target registered: %s", target_name) | ||||
|
r13196 | comm.close() | ||
|
r13195 | return | ||
self.register_comm(comm) | ||||
|
r13227 | try: | ||
f(comm, msg) | ||||
except Exception: | ||||
self.log.error("Exception opening comm with target: %s", target_name, exc_info=True) | ||||
comm.close() | ||||
self.unregister_comm(comm_id) | ||||
|
r13195 | |||
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: | ||||
# no such comm | ||||
return | ||||
|
r13227 | try: | ||
comm.handle_msg(msg) | ||||
except Exception: | ||||
self.log.error("Exception in comm_msg for %s", comm_id, exc_info=True) | ||||
|
r13195 | |||
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: | ||||
# no such comm | ||||
|
r13227 | self.log.debug("No such comm to close: %s", comm_id) | ||
|
r13195 | return | ||
del self.comms[comm_id] | ||||
|
r13227 | |||
try: | ||||
comm.handle_close(msg) | ||||
except Exception: | ||||
self.log.error("Exception handling comm_close for %s", comm_id, exc_info=True) | ||||
|
r13195 | |||
__all__ = ['CommManager'] | ||||