add-displays-to-history

Author: Darshan808

ipykernel/zmqshell.py

@@ -30,6 +30,7 @@
 from IPython.core.magic import Magics, line_magic, magics_class
 from IPython.core.magics import CodeMagics, MacroToEdit  # type:ignore[attr-defined]
 from IPython.core.usage import default_banner
+from IPython.core.history import HistoryOutput
 from IPython.display import Javascript, display
 from IPython.utils import openpy
 from IPython.utils.process import arg_split, system  # type:ignore[attr-defined]
@@ -115,6 +116,11 @@ 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'):
+            outputs = self.shell.history_manager.outputs
+            outputs[self.shell.execution_count].append(
+                HistoryOutput(output_type="display_data", bundle=data)
+            )
         self._flush_streams()
         if metadata is None:
             metadata = {}

tests/test_zmq_shell.py

@@ -209,6 +209,26 @@ def test_unregister_hook(self):
         second = self.disp_pub.unregister_hook(hook)
         assert not bool(second)
 
+    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 = {1: []}
+
+        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()