You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

91 lines
2.7 KiB
Python

"""Replacements for sys.displayhook that publish over ZMQ."""
# Copyright (c) IPython Development Team.
# Distributed under the terms of the Modified BSD License.
import builtins
import sys
from IPython.core.displayhook import DisplayHook
from jupyter_client.session import Session, extract_header
from traitlets import Any, Dict, Instance
from ipykernel.jsonutil import encode_images, json_clean
class ZMQDisplayHook:
"""A simple displayhook that publishes the object's repr over a ZeroMQ
socket."""
topic = b"execute_result"
def __init__(self, session, pub_socket):
self.session = session
self.pub_socket = pub_socket
self.parent_header = {}
def get_execution_count(self):
"""This method is replaced in kernelapp"""
return 0
def __call__(self, obj):
if obj is None:
return
builtins._ = obj # type:ignore[attr-defined]
sys.stdout.flush()
sys.stderr.flush()
contents = {
"execution_count": self.get_execution_count(),
"data": {"text/plain": repr(obj)},
"metadata": {},
}
self.session.send(
self.pub_socket, "execute_result", contents, parent=self.parent_header, ident=self.topic
)
def set_parent(self, parent):
self.parent_header = extract_header(parent)
class ZMQShellDisplayHook(DisplayHook):
"""A displayhook subclass that publishes data using ZeroMQ. This is intended
to work with an InteractiveShell instance. It sends a dict of different
representations of the object."""
topic = None
session = Instance(Session, allow_none=True)
pub_socket = Any(allow_none=True)
parent_header = Dict({})
def set_parent(self, parent):
"""Set the parent for outbound messages."""
self.parent_header = extract_header(parent)
def start_displayhook(self):
self.msg = self.session.msg(
"execute_result",
{
"data": {},
"metadata": {},
},
parent=self.parent_header,
)
def write_output_prompt(self):
"""Write the output prompt."""
self.msg["content"]["execution_count"] = self.prompt_count
def write_format_data(self, format_dict, md_dict=None):
self.msg["content"]["data"] = json_clean(encode_images(format_dict))
self.msg["content"]["metadata"] = md_dict
def finish_displayhook(self):
"""Finish up all displayhook activities."""
sys.stdout.flush()
sys.stderr.flush()
if self.msg["content"]["data"]:
self.session.send(self.pub_socket, self.msg, ident=self.topic)
self.msg = None