From 71d15f4ac0b581ea595ccb90547ec06917ffcc59 Mon Sep 17 00:00:00 2001
From: cvanelteren <caspervanelteren@gmail.com>
Date: Tue, 12 May 2026 14:46:39 +1000
Subject: [PATCH] Keep animation save layout stable after the first tight pass
 by reusing the established geometry on later dirty frames instead of
 re-running tight layout each time. Add a regression test that saves a small
 animation and verifies only the first save-time layout pass uses tight
 layout.

---
 ultraplot/figure.py               | 16 ++++++++--
 ultraplot/tests/test_animation.py | 49 +++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/ultraplot/figure.py b/ultraplot/figure.py
index 334c61a44..ab8f1824c 100644
--- a/ultraplot/figure.py
+++ b/ultraplot/figure.py
@@ -496,6 +496,12 @@ def _canvas_preprocess(self, *args, **kwargs):
 
         skip_autolayout = getattr(fig, "_skip_autolayout", False)
         layout_dirty = getattr(fig, "_layout_dirty", False)
+        saving_frame_count = getattr(fig, "_saving_frame_count", 0)
+        lock_tight_during_save = (
+            getattr(self, "_is_saving", False)
+            and saving_frame_count > 0
+            and getattr(fig, "_tight_active", False)
+        )
         if (
             skip_autolayout
             and getattr(fig, "_layout_initialized", False)
@@ -514,14 +520,20 @@ def _canvas_preprocess(self, *args, **kwargs):
         with ctx1, ctx2, ctx3:
             needs_post_layout = False
             if not fig._layout_initialized or layout_dirty:
-                fig.auto_layout()
+                fig.auto_layout(tight=False if lock_tight_during_save else None)
                 fig._layout_initialized = True
                 fig._layout_dirty = False
-                needs_post_layout = _needs_post_tight_layout(fig)
+                needs_post_layout = (
+                    not lock_tight_during_save and _needs_post_tight_layout(fig)
+                )
             result = func(self, *args, **kwargs)
             if needs_post_layout:
                 fig.auto_layout()
                 result = func(self, *args, **kwargs)
+            if method == "print_figure" and getattr(self, "_is_saving", False):
+                fig._saving_frame_count = saving_frame_count + 1
+            elif not getattr(self, "_is_saving", False):
+                fig._saving_frame_count = 0
             return result
 
     # Add preprocessor
diff --git a/ultraplot/tests/test_animation.py b/ultraplot/tests/test_animation.py
index 6e8ad2efc..00ee7c007 100644
--- a/ultraplot/tests/test_animation.py
+++ b/ultraplot/tests/test_animation.py
@@ -1,5 +1,6 @@
 from unittest.mock import MagicMock
 
+import matplotlib
 import numpy as np
 import pytest
 from matplotlib.animation import FuncAnimation
@@ -58,3 +59,51 @@ def update(frame):
     ani = FuncAnimation(fig, update, frames=10)
     # The test passes if no exception is raised
     fig.canvas.draw()
+
+
+def test_animation_save_only_tightens_first_frame(tmp_path):
+    """
+    Saving an animation should not rerun tight layout on every frame after the
+    first saved frame, or frame geometry can shift between outputs.
+    """
+    matplotlib.use("Agg")
+    state = np.random.RandomState(51423)
+
+    fig, axs = uplt.subplots(nrows=1, ncols=2, width="14cm")
+    mappables = []
+    for ax in axs:
+        m = ax.heatmap(state.rand(10, 10), cmap="dusk")
+        ax.colorbar(m, loc="t", tickdir="out", label="Axes Colorbars")
+        mappables.append(m)
+
+    axs.format(
+        abc="(a)",
+        abcloc="ul",
+        xlabel="xlabel",
+        ylabel="ylabel",
+        toplabels=("Left Axes", "Right Axes"),
+        urtitle="1",
+        suptitle="Test Animation",
+    )
+
+    auto_layout_calls = []
+    original_auto_layout = fig.auto_layout
+
+    def wrapped_auto_layout(*args, **kwargs):
+        auto_layout_calls.append(kwargs.get("tight", None))
+        return original_auto_layout(*args, **kwargs)
+
+    fig.auto_layout = wrapped_auto_layout
+
+    def update(frame):
+        for m in mappables:
+            m.set_array(state.rand(10, 10))
+        axs.format(urtitle=f"{frame + 1}")
+        return mappables
+
+    ani = FuncAnimation(fig, update, frames=3, interval=150)
+    ani.save(tmp_path / "test_animation.gif", writer="pillow")
+
+    assert auto_layout_calls
+    assert auto_layout_calls[0] is not False
+    assert auto_layout_calls[1:] == [False] * (len(auto_layout_calls) - 1)