Store display outputs in history for %notebook magic (#1435)

Author: Darshan808

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

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