diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 37575ee2d..ba707d481 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -35,12 +35,17 @@ from IPython.utils.process import arg_split, system # type:ignore[attr-defined] from jupyter_client.session import Session, extract_header from jupyter_core.paths import jupyter_runtime_dir -from traitlets import Any, CBool, CBytes, Instance, Type, default, observe +from traitlets import Any, Bool, CBool, CBytes, Instance, Type, default, observe from ipykernel import connect_qtconsole, get_connection_file, get_connection_info from ipykernel.displayhook import ZMQShellDisplayHook from ipykernel.jsonutil import encode_images, json_clean +try: + from IPython.core.history import HistoryOutput +except ImportError: + HistoryOutput = None # type: ignore[assignment,misc] + # ----------------------------------------------------------------------------- # Functions and classes # ----------------------------------------------------------------------------- @@ -54,6 +59,11 @@ class ZMQDisplayPublisher(DisplayPublisher): _parent_header: contextvars.ContextVar[dict[str, Any]] topic = CBytes(b"display_data") + store_display_history = Bool( + False, + help="If set to True, store display outputs in the history manager. Default is False.", + ).tag(config=True) + # thread_local: # An attribute used to ensure the correct output message # is processed. See ipykernel Issue 113 for a discussion. @@ -115,6 +125,21 @@ def publish( # type:ignore[override] update : bool, optional, keyword-only If True, send an update_display_data message instead of display_data. """ + if ( + self.store_display_history + and self.shell is not None + and hasattr(self.shell, "history_manager") + and HistoryOutput is not None + ): + # Reference: github.com/ipython/ipython/pull/14998 + exec_count = self.shell.execution_count + if getattr(self.shell.display_pub, "_in_post_execute", False): + exec_count -= 1 + outputs = getattr(self.shell.history_manager, "outputs", None) + if outputs is not None: + outputs.setdefault(exec_count, []).append( + HistoryOutput(output_type="display_data", bundle=data) + ) self._flush_streams() if metadata is None: metadata = {} diff --git a/tests/test_zmq_shell.py b/tests/test_zmq_shell.py index bc5e3f556..c4859cfa0 100644 --- a/tests/test_zmq_shell.py +++ b/tests/test_zmq_shell.py @@ -22,6 +22,11 @@ ZMQInteractiveShell, ) +try: + from IPython.core.history import HistoryOutput +except ImportError: + HistoryOutput = None # type: ignore[assignment,misc] + class NoReturnDisplayHook: """ @@ -209,6 +214,35 @@ def test_unregister_hook(self): second = self.disp_pub.unregister_hook(hook) assert not bool(second) + @unittest.skipIf(HistoryOutput is None, "HistoryOutput not available") + def test_display_stored_in_history(self): + """ + Test that published display data gets stored in shell history + for %notebook magic support, and not stored when disabled. + """ + for enable in [False, True]: + # Mock shell with history manager + mock_shell = MagicMock() + mock_shell.execution_count = 1 + mock_shell.history_manager.outputs = dict() + mock_shell.display_pub._in_post_execute = False + + self.disp_pub.shell = mock_shell + self.disp_pub.store_display_history = enable + + data = {"text/plain": "test output"} + self.disp_pub.publish(data) + + if enable: + # Check that output was stored in history + stored_outputs = mock_shell.history_manager.outputs[1] + assert len(stored_outputs) == 1 + assert stored_outputs[0].output_type == "display_data" + assert stored_outputs[0].bundle == data + else: + # Should not store anything in history + assert mock_shell.history_manager.outputs == {} + def test_magics(tmp_path): context = zmq.Context()