fix: prevent mpl figure DPI from compounding on cell rerun (#9474)

Author: kirangadhave

marimo/_plugins/stateless/mpl/_mpl.py

@@ -128,6 +128,17 @@ def interactive(figure: Figure | SubFigure | Axes) -> Html:
     if not ctx.virtual_files_supported:
         return NonInteractiveMplHtml(figure)
 
+    # Figure::figure returns self; SubFigure::figure returns the parent Figure
+    is_subfigure = figure.figure is not figure
+    if is_subfigure:
+        warnings.warn(
+            message="SubFigure is not supported in interactive mode; "
+            "rendering as static PNG instead. "
+            "Consider using a regular Figure instead.",
+            stacklevel=2,
+        )
+        return NonInteractiveMplHtml(figure=figure)
+
     from marimo._plugins.ui._impl.from_mpl_interactive import mpl_interactive
 
     return mpl_interactive(figure)

marimo/_plugins/ui/_impl/from_mpl_interactive.py

@@ -163,6 +163,12 @@ class mpl_interactive(UIElement[ModelIdRef, dict[str, Any]]):
     def __init__(self, figure: Figure | SubFigure) -> None:
         self._figure = figure
 
+        # SubFigure delegates dpi/size_inches to its parent Figure;
+        # Figure.figure returns self, so this works for both.
+        root = figure.figure
+        self._original_dpi = root.get_dpi()
+        self._original_size_inches = tuple(root.get_size_inches())
+
         # Create FigureManagerWebAgg
         self._figure_manager = new_figure_manager_given_figure(
             id(figure), figure
@@ -233,7 +239,11 @@ def _initialize(
             ctx = get_context()
             ctx.cell_lifecycle_registry.add(
                 _MplCleanupHandle(
-                    self._comm, self._figure_manager, self._sync_ws
+                    comm=self._comm,
+                    figure_manager=self._figure_manager,
+                    sync_ws=self._sync_ws,
+                    original_dpi=self._original_dpi,
+                    original_size_inches=self._original_size_inches,
                 )
             )
         except ContextNotInitializedError:
@@ -302,12 +312,16 @@ class _MplCleanupHandle(CellLifecycleItem):
     def __init__(
         self,
         comm: MarimoComm,
+        original_dpi: float,
+        original_size_inches: tuple[float, float],
         figure_manager: Any = None,
         sync_ws: Any = None,
     ) -> None:
         self._comm = comm
         self._figure_manager = figure_manager
         self._sync_ws = sync_ws
+        self._original_dpi = original_dpi
+        self._original_size_inches = original_size_inches
 
     def create(self, context: RuntimeContext) -> None:
         del context
@@ -323,9 +337,18 @@ def dispose(self, context: RuntimeContext, deletion: bool) -> bool:
             except Exception:
                 pass
         if self._figure_manager is not None:
+            try:
+                # get the root figure (in case of Subfigure) which handles dpi
+                root = self._figure_manager.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
+
         self._comm.close()
         return True

tests/_plugins/stateless/test_mpl.py

@@ -147,6 +147,48 @@ def test_mpl_interactive_fallback_when_virtual_files_not_supported() -> None:
     plt.close(fig)
 
 
+@pytest.mark.requires("matplotlib")
+def test_mpl_interactive_subfigure_falls_back_to_png() -> None:
+    """SubFigure cannot be attached to a WebAgg canvas (matplotlib limitation),
+    so mo.mpl.interactive should fall back to a static PNG and warn.
+    """
+    import warnings
+
+    import matplotlib.pyplot as plt
+
+    from marimo._plugins.stateless.mpl._mpl import (
+        NonInteractiveMplHtml,
+        interactive,
+    )
+    from marimo._runtime.context.kernel_context import KernelRuntimeContext
+
+    parent = plt.figure(figsize=(8, 4), dpi=100)
+    sub_left, _sub_right = parent.subfigures(1, 2)
+    sub_left.subplots().plot([1, 2, 3])
+
+    mock_ctx = MagicMock(spec=KernelRuntimeContext)
+    mock_ctx.virtual_files_supported = True
+
+    with (
+        patch(
+            "marimo._plugins.stateless.mpl._mpl.get_context",
+            return_value=mock_ctx,
+        ),
+        warnings.catch_warnings(record=True) as captured,
+    ):
+        warnings.simplefilter("always")
+        result = interactive(sub_left)
+
+    assert isinstance(result, NonInteractiveMplHtml)
+
+    # A UserWarning explaining the fallback should have been emitted.
+    subfigure_warnings = [w for w in captured if "SubFigure" in str(w.message)]
+    assert len(subfigure_warnings) == 1
+    assert issubclass(subfigure_warnings[0].category, UserWarning)
+
+    plt.close(parent)
+
+
 @pytest.mark.requires("matplotlib")
 def test_new_figure_manager_suppresses_thread_warning() -> None:
     """Regression test for https://github.com/marimo-team/marimo/issues/8747.

tests/_plugins/ui/_impl/test_from_mpl_interactive.py

@@ -6,6 +6,11 @@
 
 import pytest
 
+mpl = pytest.importorskip("matplotlib")
+mpl.use(
+    "Agg"
+)  # Non-interactive backend; avoids DPR pre-inflation on HiDPI hosts.
+
 
 @pytest.mark.requires("matplotlib")
 class TestSyncWebSocket:
@@ -180,6 +185,63 @@ def capture_send(data: Any, buffers: Any = None) -> None:
         plt.close(fig)
 
 
+@pytest.mark.requires("matplotlib")
+class TestDpiPreservationOnRerun:
+    """Re-running a cell that wraps the same figure should not compound DPI.
+
+    matplotlib's ``FigureCanvasBase.__init__`` unconditionally captures
+    ``figure._original_dpi = figure.dpi``. After a HiDPI client connects
+    and scales ``figure.dpi`` up by the device pixel ratio, a subsequent
+    canvas creation on the same figure would treat that scaled value as
+    "original" and scale it again — making the resolution compound on
+    every rerun (see issue #9466).
+    """
+
+    def test_dpi_does_not_compound_across_reruns(self) -> None:
+        import matplotlib.pyplot as plt
+
+        from marimo._plugins.ui._impl.from_mpl_interactive import (
+            _MplCleanupHandle,
+            mpl_interactive,
+        )
+
+        fig, ax = plt.subplots(figsize=(5, 5), dpi=100)
+        ax.plot([1, 2, 3])
+
+        for _ in range(3):
+            with patch("marimo._plugins.ui._impl.comm.broadcast_notification"):
+                element = mpl_interactive(fig)
+
+            # Simulate the HiDPI handshake from the frontend.
+            element._figure_manager.handle_json(
+                {"type": "set_device_pixel_ratio", "device_pixel_ratio": 2}
+            )
+            element._figure_manager.handle_json(
+                {"type": "resize", "width": 500, "height": 500}
+            )
+            # While the canvas is live, dpi reflects the device-scaled value.
+            assert fig.dpi == 200
+            assert tuple(fig.get_size_inches()) == (5.0, 5.0)
+
+            # Simulate cell teardown — the cleanup handle is what marimo
+            # registers via cell_lifecycle_registry; running it directly
+            # avoids needing a live runtime context in the test.
+            cleanup = _MplCleanupHandle(
+                comm=element._comm,
+                figure_manager=element._figure_manager,
+                sync_ws=element._sync_ws,
+                original_dpi=element._original_dpi,
+                original_size_inches=element._original_size_inches,
+            )
+            cleanup.dispose(context=MagicMock(), deletion=False)
+
+            # After dispose the figure is restored to the user's intent.
+            assert fig.dpi == 100
+            assert tuple(fig.get_size_inches()) == (5.0, 5.0)
+
+        plt.close(fig)
+
+
 @pytest.mark.requires("matplotlib")
 class TestMplCleanupHandle:
     """Test that _MplCleanupHandle properly closes the comm."""
@@ -190,7 +252,9 @@ def test_dispose_closes_comm(self) -> None:
         )
 
         mock_comm = MagicMock()
-        handle = _MplCleanupHandle(mock_comm)
+        handle = _MplCleanupHandle(
+            mock_comm, original_dpi=100, original_size_inches=(5.0, 5.0)
+        )
         result = handle.dispose(context=MagicMock(), deletion=False)
 
         assert result is True
@@ -202,7 +266,9 @@ def test_dispose_on_deletion(self) -> None:
         )
 
         mock_comm = MagicMock()
-        handle = _MplCleanupHandle(mock_comm)
+        handle = _MplCleanupHandle(
+            mock_comm, original_dpi=100, original_size_inches=(5.0, 5.0)
+        )
         result = handle.dispose(context=MagicMock(), deletion=True)
 
         assert result is True
@@ -216,7 +282,13 @@ def test_dispose_cleans_up_figure_manager(self) -> None:
         mock_comm = MagicMock()
         mock_manager = MagicMock()
         mock_ws = MagicMock()
-        handle = _MplCleanupHandle(mock_comm, mock_manager, mock_ws)
+        handle = _MplCleanupHandle(
+            mock_comm,
+            figure_manager=mock_manager,
+            sync_ws=mock_ws,
+            original_dpi=100,
+            original_size_inches=(5.0, 5.0),
+        )
         result = handle.dispose(context=MagicMock(), deletion=False)
 
         assert result is True
@@ -234,7 +306,13 @@ def test_dispose_tolerates_manager_errors(self) -> None:
         mock_manager.remove_web_socket.side_effect = RuntimeError("boom")
         mock_manager.canvas.close.side_effect = RuntimeError("boom")
         mock_ws = MagicMock()
-        handle = _MplCleanupHandle(mock_comm, mock_manager, mock_ws)
+        handle = _MplCleanupHandle(
+            mock_comm,
+            figure_manager=mock_manager,
+            sync_ws=mock_ws,
+            original_dpi=100,
+            original_size_inches=(5.0, 5.0),
+        )
         # Should not raise even if manager cleanup fails
         result = handle.dispose(context=MagicMock(), deletion=False)