1  """Replacements for sys.displayhook that publish over ZMQ."""
   2  
   3  # Copyright (c) IPython Development Team.
   4  # Distributed under the terms of the Modified BSD License.
   5  from __future__ import annotations
   6  
   7  import builtins
   8  import sys
   9  import typing as t
  10  
  11  from IPython.core.displayhook import DisplayHook
  12  from jupyter_client.session import Session, extract_header
  13  from traitlets import Any, Dict, Instance
  14  
  15  from ipykernel.jsonutil import encode_images, json_clean
  16  
  17  
  18  class ZMQDisplayHook:
  19      """A simple displayhook that publishes the object's repr over a ZeroMQ
  20      socket."""
  21  
  22      topic = b"execute_result"
  23  
  24      def __init__(self, session, pub_socket):
  25          """Initialize the hook."""
  26          self.session = session
  27          self.pub_socket = pub_socket
  28          self.parent_header = {}
  29  
  30      def get_execution_count(self):
  31          """This method is replaced in kernelapp"""
  32          return 0
  33  
  34      def __call__(self, obj):
  35          """Handle a hook call."""
  36          if obj is None:
  37              return
  38  
  39          builtins._ = obj  # type:ignore[attr-defined]
  40          sys.stdout.flush()
  41          sys.stderr.flush()
  42          contents = {
  43              "execution_count": self.get_execution_count(),
  44              "data": {"text/plain": repr(obj)},
  45              "metadata": {},
  46          }
  47          self.session.send(
  48              self.pub_socket, "execute_result", contents, parent=self.parent_header, ident=self.topic
  49          )
  50  
  51      def set_parent(self, parent):
  52          """Set the parent header."""
  53          self.parent_header = extract_header(parent)
  54  
  55  
  56  class ZMQShellDisplayHook(DisplayHook):
  57      """A displayhook subclass that publishes data using ZeroMQ. This is intended
  58      to work with an InteractiveShell instance. It sends a dict of different
  59      representations of the object."""
  60  
  61      topic = None
  62  
  63      session = Instance(Session, allow_none=True)
  64      pub_socket = Any(allow_none=True)
  65      parent_header = Dict({})
  66      msg: dict[str, t.Any] | None
  67  
  68      def set_parent(self, parent):
  69          """Set the parent for outbound messages."""
  70          self.parent_header = extract_header(parent)
  71  
  72      def start_displayhook(self):
  73          """Start the display hook."""
  74          if self.session:
  75              self.msg = self.session.msg(
  76                  "execute_result",
  77                  {
  78                      "data": {},
  79                      "metadata": {},
  80                  },
  81                  parent=self.parent_header,
  82              )
  83  
  84      def write_output_prompt(self):
  85          """Write the output prompt."""
  86          if self.msg:
  87              self.msg["content"]["execution_count"] = self.prompt_count
  88  
  89      def write_format_data(self, format_dict, md_dict=None):
  90          """Write format data to the message."""
  91          if self.msg:
  92              self.msg["content"]["data"] = json_clean(encode_images(format_dict))
  93              self.msg["content"]["metadata"] = md_dict
  94  
  95      def finish_displayhook(self):
  96          """Finish up all displayhook activities."""
  97          sys.stdout.flush()
  98          sys.stderr.flush()
  99          if self.msg and self.msg["content"]["data"] and self.session:
 100              self.session.send(self.pub_socket, self.msg, ident=self.topic)
 101          self.msg = None
 102