fix(mpl): disconnect toolbar callbacks on cell rerun

marimo/_plugins/ui/_impl/from_mpl_interactive.py

@@ -29,6 +29,31 @@
 
 LOGGER = _loggers.marimo_logger()
 
+
+def _disconnect_owner_callbacks(canvas: Any, owner: Any) -> None:
+    """Disconnect every callback on ``canvas.callbacks`` whose receiver is ``owner``.
+
+    ``FigureCanvasBase.callbacks`` is a property returning
+    ``figure._canvas_callbacks``, so every canvas bound to the same
+    figure shares one registry. ``NavigationToolbar2.__init__``
+    registers ``_zoom_pan_handler`` (and a few siblings) on this
+    shared registry; matplotlib has no public counterpart that
+    disconnects them when a canvas/toolbar is discarded. Without this
+    cleanup, every cell rerun would stack another toolbar's handler on
+    the registry, leading to duplicate ``press_pan``/``release_pan``
+    dispatch.
+    """
+    registry = canvas.callbacks
+    cids: list[int] = []
+    for handlers in registry.callbacks.values():
+        for cid, ref in handlers.items():
+            fn = ref() if callable(ref) else ref
+            if getattr(fn, "__self__", None) is owner:
+                cids.append(cid)
+    for cid in cids:
+        registry.disconnect(cid)
+
+
 # Must match the className on the container div in MplInteractivePlugin.tsx
 _MPL_SCOPE = ".mpl-interactive-figure"
 
@@ -328,27 +353,33 @@ def create(self, context: RuntimeContext) -> None:
 
     def dispose(self, context: RuntimeContext, deletion: bool) -> bool:
         del context, deletion
-        # Disconnect the web socket from the figure manager first,
-        # then destroy the canvas to prevent stale timer callbacks
-        # on thread teardown (avoids "QObject::killTimer" errors).
-        if self._figure_manager is not None and self._sync_ws is not None:
-            try:
-                self._figure_manager.remove_web_socket(self._sync_ws)
-            except Exception:
-                pass
-        if self._figure_manager is not None:
+        fm = self._figure_manager
+        if fm is not None:
+            if self._sync_ws is not None:
+                try:
+                    fm.remove_web_socket(self._sync_ws)
+                except Exception:
+                    LOGGER.exception("Failed to remove mpl web socket")
+
+            canvas = fm.canvas
+            toolbar = getattr(canvas, "toolbar", None)
+            if toolbar is not None:
+                try:
+                    _disconnect_owner_callbacks(canvas, toolbar)
+                except Exception:
+                    LOGGER.exception(
+                        "Failed to disconnect mpl toolbar callbacks"
+                    )
+
             try:
                 # get the root figure (in case of Subfigure) which handles dpi
-                root = self._figure_manager.canvas.figure.figure
+                root = canvas.figure.figure
                 root.set_dpi(self._original_dpi)
                 root.set_size_inches(*self._original_size_inches)
             except Exception:
-                pass
-
-            try:
-                self._figure_manager.canvas.close()
-            except Exception:
-                pass
+                LOGGER.exception(
+                    "Failed to restore mpl figure dpi/size on dispose"
+                )
 
         self._comm.close()
         return True

tests/_plugins/ui/_impl/test_from_mpl_interactive.py

@@ -281,6 +281,9 @@ def test_dispose_cleans_up_figure_manager(self) -> None:
 
         mock_comm = MagicMock()
         mock_manager = MagicMock()
+        # The disconnect helper iterates this; keep it empty so the mock
+        # doesn't trip over MagicMock's default value generators.
+        mock_manager.canvas.callbacks.callbacks = {}
         mock_ws = MagicMock()
         handle = _MplCleanupHandle(
             mock_comm,
@@ -293,7 +296,6 @@ def test_dispose_cleans_up_figure_manager(self) -> None:
 
         assert result is True
         mock_manager.remove_web_socket.assert_called_once_with(mock_ws)
-        mock_manager.canvas.close.assert_called_once()
         mock_comm.close.assert_called_once()
 
     def test_dispose_tolerates_manager_errors(self) -> None:
@@ -304,7 +306,10 @@ def test_dispose_tolerates_manager_errors(self) -> None:
         mock_comm = MagicMock()
         mock_manager = MagicMock()
         mock_manager.remove_web_socket.side_effect = RuntimeError("boom")
-        mock_manager.canvas.close.side_effect = RuntimeError("boom")
+        # Force the disconnect path to blow up too.
+        type(mock_manager.canvas).callbacks = property(
+            lambda _self: (_ for _ in ()).throw(RuntimeError("boom"))
+        )
         mock_ws = MagicMock()
         handle = _MplCleanupHandle(
             mock_comm,
@@ -320,6 +325,73 @@ def test_dispose_tolerates_manager_errors(self) -> None:
         mock_comm.close.assert_called_once()
 
 
+@pytest.mark.requires("matplotlib")
+class TestToolbarCallbackCleanup:
+    """Disposing a cleanup handle must disconnect the toolbar's callbacks
+    from the figure-shared event registry.
+
+    `canvas.callbacks` is a property that delegates to
+    `figure._canvas_callbacks`, so every canvas bound to the same figure
+    shares one registry. Without explicit cleanup, each cell rerun stacks
+    another `_zoom_pan_handler` on the registry — leading to duplicate
+    `press_pan`/`release_pan` dispatch and ultimately the
+    `AttributeError: 'Axes' object has no attribute '_pan_start'` from
+    matplotlib's `Axes.end_pan`.
+    """
+
+    @staticmethod
+    def _count_toolbar_handlers(fig: Any, signal: str) -> int:
+        handlers = fig.canvas.callbacks.callbacks.get(signal, {})
+        count = 0
+        for ref in handlers.values():
+            fn = ref() if callable(ref) else ref
+            owner = getattr(fn, "__self__", None)
+            if (
+                owner is not None
+                and type(owner).__name__ == "NavigationToolbar2WebAgg"
+            ):
+                count += 1
+        return count
+
+    def test_dispose_disconnects_toolbar_callbacks(self) -> None:
+        import matplotlib.pyplot as plt
+
+        from marimo._plugins.ui._impl.from_mpl_interactive import (
+            _MplCleanupHandle,
+            mpl_interactive,
+        )
+
+        fig, ax = plt.subplots()
+        ax.plot([1, 2, 3])
+
+        with patch("marimo._plugins.ui._impl.comm.broadcast_notification"):
+            elem1 = mpl_interactive(fig)
+        assert self._count_toolbar_handlers(fig, "button_press_event") == 1
+        assert self._count_toolbar_handlers(fig, "button_release_event") == 1
+
+        # Simulate cell teardown: marimo's runtime invokes dispose on the
+        # lifecycle handle when the cell re-runs or is deleted.
+        cleanup = _MplCleanupHandle(
+            comm=elem1._comm,
+            figure_manager=elem1._figure_manager,
+            sync_ws=elem1._sync_ws,
+            original_dpi=elem1._original_dpi,
+            original_size_inches=elem1._original_size_inches,
+        )
+        cleanup.dispose(context=MagicMock(), deletion=False)
+
+        # Simulate the cell re-running: a new mpl_interactive on the same
+        # figure object. After dispose, only the new toolbar's callbacks
+        # should be live on the shared registry.
+        with patch("marimo._plugins.ui._impl.comm.broadcast_notification"):
+            _elem2 = mpl_interactive(fig)
+
+        assert self._count_toolbar_handlers(fig, "button_press_event") == 1
+        assert self._count_toolbar_handlers(fig, "button_release_event") == 1
+
+        plt.close(fig)
+
+
 @pytest.mark.requires("matplotlib")
 class TestMplInteractiveArgs:
     """Test that mpl_interactive passes correct args to the UIElement."""