diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 37575ee2..41c88318 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -41,6 +41,11 @@ 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 # ----------------------------------------------------------------------------- @@ -115,6 +120,20 @@ def publish( # type:ignore[override] update : bool, optional, keyword-only If True, send an update_display_data message instead of display_data. """ + if ( + 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 bc5e3f55..6e0f9048 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,29 @@ 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. + """ + # 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 + + data = {"text/plain": "test output"} + self.disp_pub.publish(data) + + # 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 + def test_magics(tmp_path): context = zmq.Context()