Backport PR matplotlib#31304: Fix restoring 'auto' aspect in 3D axes after switching from 'equal'

Author: jayaprajapatii

lib/mpl_toolkits/mplot3d/axes3d.py

@@ -339,9 +339,28 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
                            aspect=aspect)
 
         self.set_adjustable(adjustable)
+
+        if self._aspect == 'auto' and aspect != 'auto':
+            self._auto_aspect_restore_limits = list(self.get_w_lims())
+
         self._aspect = aspect
 
-        if aspect in ('equal', 'equalxy', 'equalxz', 'equalyz'):
+        if aspect == "auto":
+            if self._adjustable == "box":
+                self.set_box_aspect(None)
+
+            elif self._adjustable == "datalim":
+                if hasattr(self, "_auto_aspect_restore_limits"):
+                    limits = self._auto_aspect_restore_limits
+
+                    if limits[0] is not None:
+                        self.set_xlim3d(limits[0], limits[1], auto=True)
+                    if limits[2] is not None:
+                        self.set_ylim3d(limits[2], limits[3], auto=True)
+                    if limits[4] is not None:
+                        self.set_zlim3d(limits[4], limits[5], auto=True)
+                    del self._auto_aspect_restore_limits
+        else:
             ax_indices = self._equal_aspect_axis_indices(aspect)
 
             view_intervals = np.array([self.xaxis.get_view_interval(),
@@ -871,6 +890,11 @@ def _set_lim3d(self, axis, lower=None, upper=None, *, emit=True,
             delta = (upper - lower) * view_margin
             lower -= delta
             upper += delta
+
+        if not auto and hasattr(self, '_auto_aspect_restore_limits'):
+            idx = {self.xaxis: 0, self.yaxis: 2, self.zaxis: 4}[axis]
+            self._auto_aspect_restore_limits[idx] = None
+            self._auto_aspect_restore_limits[idx + 1] = None
         return axis._set_lim(lower, upper, emit=emit, auto=auto)
 
     def set_xlim(self, left=None, right=None, *, emit=True, auto=False,

lib/mpl_toolkits/mplot3d/tests/test_axes3d.py

@@ -98,6 +98,47 @@ def test_axis_positions_inverted():
         ax.set(xlabel='x', ylabel='y', zlabel='z', title=title)
 
 
+def test_set_aspect_datalim_restores_limits():
+
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    ax.plot([0,1], [0,2], [0,3])
+
+    fig.canvas.draw()
+    default_limits = ax.get_w_lims()
+
+    ax.set_aspect('equal', adjustable='datalim')
+    fig.canvas.draw()
+    equal_limits = ax.get_w_lims()
+
+    ax.set_aspect('auto', adjustable='datalim')
+    fig.canvas.draw()
+    final_limits = ax.get_w_lims()
+
+    # equal should change limits
+    assert not np.allclose(default_limits, equal_limits)
+
+    # auto should restore original limits
+    assert np.allclose(default_limits, final_limits)
+
+
+def test_set_aspect_datalim_restores_untouched_axes():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+    ax.plot([1, 2], [3, 4], [5, 6])
+    orig_lims = ax.get_w_lims()
+
+    ax.set_aspect('equal', adjustable='datalim')
+    ax.set_xlim(-50, 50)  # explicit user change mid-'equal'
+    ax.set_aspect('auto', adjustable='datalim')
+
+    xlim, ylim, zlim = ax.get_xlim3d(), ax.get_ylim3d(), ax.get_zlim3d()
+    assert xlim == (-50, 50)               # explicit value preserved
+    assert ylim == orig_lims[2:4]          # untouched axis restored
+    assert zlim == orig_lims[4:6]          # untouched axis restored
+
+
 @mpl3d_image_comparison(['aspects.png'], remove_text=False, style='mpl20')
 def test_aspects():
     aspects = ('auto', 'equal', 'equalxy', 'equalyz', 'equalxz', 'equal')