Store display outputs in history for %notebook magic

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 = {}

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()