Fix matplotlib plots displaying in wrong cells during %notebook export (ipython#14998)

## Fixes ipython#14990

**Problem:** When using `%notebook` magic to export IPython sessions,
auto flushed matplotlib plots appeared one cell later than where they
were created. This happened because matplotlib's `flush_figures()` runs
during the `post_execute` event, after the cell's output capture phase.

**Solution:** Added execution state tracking to the display publisher:
- Track when we're in `post_execute` phase vs normal execution
- When in `post_execute`: use `execution_count - 1` (associate with
originating cell)
- When in normal execution: use current `execution_count` (for explicit
`plt.show()`)

**Result:** 
- Plots without `plt.show()` now appear in the correct cell
- Plots with explicit `plt.show()` continue to work correctly
- No breaking changes to existing behavior

**Testing:** Added test to verify both auto-flush and manual-flush
scenarios place plots in the correct cells.

Author: krassowski

IPython/core/displaypub.py

@@ -43,6 +43,9 @@ class DisplayPublisher(Configurable):
     def __init__(self, shell=None, *args, **kwargs):
         self.shell = shell
         self._is_publishing = False
+        self._in_post_execute = False
+        if self.shell:
+            self._setup_execution_tracking()
         super().__init__(*args, **kwargs)
 
     def _validate_data(self, data, metadata=None):
@@ -62,6 +65,19 @@ def _validate_data(self, data, metadata=None):
             if not isinstance(metadata, dict):
                 raise TypeError("metadata must be a dict, got: %r" % data)
 
+    def _setup_execution_tracking(self):
+        """Set up hooks to track execution state"""
+        self.shell.events.register("post_execute", self._on_post_execute)
+        self.shell.events.register("pre_execute", self._on_pre_execute)
+
+    def _on_post_execute(self):
+        """Called at start of post_execute phase"""
+        self._in_post_execute = True
+
+    def _on_pre_execute(self):
+        """Called at start of pre_execute phase"""
+        self._in_post_execute = False
+
     # use * to indicate transient, update are keyword-only
     def publish(
         self,
@@ -133,7 +149,13 @@ def publish(
 
         outputs = self.shell.history_manager.outputs
 
-        outputs[self.shell.execution_count].append(
+        target_execution_count = self.shell.execution_count
+        if self._in_post_execute:
+            # We're in post_execute, so this is likely a matplotlib flush
+            # Use execution_count - 1 to associate with the cell that created the plot
+            target_execution_count = self.shell.execution_count - 1
+
+        outputs[target_execution_count].append(
             HistoryOutput(output_type="display_data", bundle=data)
         )
 

tests/test_display_2.py

@@ -201,15 +201,77 @@ def test_set_matplotlib_formats_kwargs():
     cfg = _get_inline_config()
     cfg.print_figure_kwargs.update(dict(foo="bar"))
     kwargs = dict(dpi=150)
-    set_matplotlib_formats("png", **kwargs)
-    formatter = ip.display_formatter.formatters["image/png"]
-    f = formatter.lookup_by_type(Figure)
-    formatter_kwargs = f.keywords
-    expected = kwargs
-    expected["base64"] = True
-    expected["fmt"] = "png"
-    expected.update(cfg.print_figure_kwargs)
-    assert formatter_kwargs == expected
+    try:
+        set_matplotlib_formats("png", **kwargs)
+        formatter = ip.display_formatter.formatters["image/png"]
+        f = formatter.lookup_by_type(Figure)
+        formatter_kwargs = f.keywords
+        expected = kwargs
+        expected["base64"] = True
+        expected["fmt"] = "png"
+        expected.update(cfg.print_figure_kwargs)
+        assert formatter_kwargs == expected
+    finally:
+        cfg.print_figure_kwargs.clear()
+
+
+@dec.skip_without("matplotlib")
+def test_matplotlib_positioning():
+    _ip = get_ipython()
+
+    prev_active_types = _ip.display_formatter.active_types.copy()
+    prev_execution_count = _ip.execution_count
+    prev_user_ns_underscore = _ip.user_ns.get("_", None)
+
+    _ip.history_manager.reset()
+    _ip.display_formatter.active_types = ["text/plain", "image/png"]
+
+    _ip.run_cell("import matplotlib")
+    prev_mpl_backend = _ip.run_cell("matplotlib.get_backend()").result
+
+    try:
+        _ip.run_line_magic("matplotlib", "inline")
+        _ip.execution_count = 1
+        _ip.run_cell("'no plot here'", store_history=True)
+
+        # Cell 2: No manual flush
+        _ip.run_cell(
+            "import matplotlib.pyplot as plt;plt.plot([0, 1])", store_history=True
+        )
+
+        _ip.run_cell("'no plot here'", store_history=True)
+
+        # Cell 4: Manual flush
+        _ip.run_cell("plt.plot([1, 0])\nplt.show()", store_history=True)
+
+        _ip.run_cell("'no plot here'", store_history=True)
+
+        outputs = _ip.history_manager.outputs
+
+        # Only cells 2 and 4 should have plots
+        for cell_num in [1, 3, 5]:
+            assert not any(
+                "image/png" in out.bundle for out in outputs.get(cell_num, [])
+            ), f"Cell {cell_num} should not have plot"
+
+        cell_2_has_plot = any("image/png" in out.bundle for out in outputs.get(2, []))
+        cell_4_has_plot = any("image/png" in out.bundle for out in outputs.get(4, []))
+
+        assert cell_2_has_plot, "Cell 2 should have plot (auto-flush)"
+        assert cell_4_has_plot, "Cell 4 should have plot (manual flush)"
+
+    finally:
+        _ip.run_cell("plt.close('all')")
+        _ip.run_line_magic("matplotlib", prev_mpl_backend)
+        _ip.history_manager.reset()
+        _ip.display_formatter.active_types = prev_active_types
+        _ip.displayhook.flush()
+
+        _ip.execution_count = prev_execution_count
+        if prev_user_ns_underscore is not None:
+            _ip.user_ns["_"] = prev_user_ns_underscore
+        elif "_" in _ip.user_ns:
+            del _ip.user_ns["_"]
 
 
 def test_display_available():