Restore _AxesStack to track a Figure's Axes order.

QuLogic · QuLogic · commit ac394fc3e192 · 2021-03-02T19:30:37.000-05:00
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -622,9 +622,6 @@ def __len__(self):
     def __getitem__(self, ind):
         return self._elements[ind]
 
-    def as_list(self):
-        return list(self._elements)
-
     def forward(self):
         """Move the position forward and return the current element."""
         self._pos = min(self._pos + 1, len(self._elements) - 1)
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -52,6 +52,71 @@ def _stale_figure_callback(self, val):
         self.figure.stale = val
 
 
+class _AxesStack(cbook.Stack):
+    """
+    Specialization of Stack, to handle all tracking of Axes in a Figure.
+
+    This stack stores ``ind, axes`` pairs, where ``ind`` is a serial index
+    tracking the order in which axes were added.
+
+    AxesStack is a callable; calling it returns the current axes.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._ind = 0
+
+    def as_list(self):
+        """
+        Return a list of the Axes instances that have been added to the figure.
+        """
+        return [a for i, a in sorted(self._elements)]
+
+    def _entry_from_axes(self, e):
+        return next(((ind, a) for ind, a in self._elements if a == e), None)
+
+    def remove(self, a):
+        """Remove the axes from the stack."""
+        super().remove(self._entry_from_axes(a))
+
+    def bubble(self, a):
+        """
+        Move the given axes, which must already exist in the stack, to the top.
+        """
+        return super().bubble(self._entry_from_axes(a))
+
+    def add(self, a):
+        """
+        Add Axes *a* to the stack, and return the stack.
+
+        If *a* is already on the stack, don't add it again, but
+        return *None*.
+        """
+        # All the error checking may be unnecessary; but this method
+        # is called so seldom that the overhead is negligible.
+        _api.check_isinstance(Axes, a=a)
+
+        if a in self:
+            return None
+        self._ind += 1
+        return super().push((self._ind, a))
+
+    def __call__(self):
+        """
+        Return the active axes.
+
+        If no axes exists on the stack, then returns None.
+        """
+        if not len(self._elements):
+            return self._default
+        else:
+            index, axes = self._elements[self._pos]
+            return axes
+
+    def __contains__(self, a):
+        return a in self.as_list()
+
+
 class SubplotParams:
     """
     A class to hold the parameters for a subplot.
@@ -141,7 +206,7 @@ def __init__(self):
         self.figure = self
         # list of child gridspecs for this figure
         self._gridspecs = []
-        self._localaxes = cbook.Stack()  # keep track of axes at this level
+        self._localaxes = _AxesStack()  # track all axes and current axes
         self.artists = []
         self.lines = []
         self.patches = []
@@ -716,8 +781,8 @@ def add_subplot(self, *args, **kwargs):
 
     def _add_axes_internal(self, ax, key):
         """Private helper for `add_axes` and `add_subplot`."""
-        self._axstack.push(ax)
-        self._localaxes.push(ax)
+        self._axstack.add(ax)
+        self._localaxes.add(ax)
         self.sca(ax)
         ax._remove_method = self.delaxes
         # this is to support plt.subplot's re-selection logic
@@ -2155,7 +2220,7 @@ def __init__(self,
 
         self.set_tight_layout(tight_layout)
 
-        self._axstack = cbook.Stack()  # track all figure axes and current axes
+        self._axstack = _AxesStack()  # track all figure axes and current axes
         self.clf()
         self._cachedRenderer = None
 
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -195,6 +195,9 @@ def test_gca():
         assert fig.gca(projection='rectilinear') is ax1
     assert fig.gca() is ax1
 
+    # sca() should not change stored order of Axes, which is order added.
+    assert fig.axes == [ax0, ax1, ax2, ax3]
+
 
 def test_add_subplot_subclass():
     fig = plt.figure()
