Backport PR #19625: Restore _AxesStack to track a Figure's Axes order.

jklymak · meeseeksmachine · commit fc362fcf29a7 · 2021-03-05T23:25:21.000Z
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -620,9 +620,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.
+
+        If *a* is already on the stack, don't add it again.
+        """
+        # 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
+
+        self._ind += 1
+        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 None
+        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
@@ -2161,7 +2226,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
@@ -188,13 +188,28 @@ def test_gca():
         assert fig.gca(polar=True) is not ax2
     assert fig.gca().get_subplotspec().get_geometry() == (1, 2, 1, 1)
 
+    # add_axes on an existing Axes should not change stored order, but will
+    # make it current.
+    fig.add_axes(ax0)
+    assert fig.axes == [ax0, ax1, ax2, ax3]
+    assert fig.gca() is ax0
+
+    # add_subplot on an existing Axes should not change stored order, but will
+    # make it current.
+    fig.add_subplot(ax2)
+    assert fig.axes == [ax0, ax1, ax2, ax3]
+    assert fig.gca() is ax2
+
     fig.sca(ax1)
     with pytest.warns(
             MatplotlibDeprecationWarning,
             match=r'Calling gca\(\) with keyword arguments was deprecated'):
         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()
@@ -241,6 +256,11 @@ def test_add_subplot_invalid():
                       match='Passing non-integers as three-element position '
                             'specification is deprecated'):
         fig.add_subplot(2.0, 2, 1)
+    _, ax = plt.subplots()
+    with pytest.raises(ValueError,
+                       match='The Subplot must have been created in the '
+                             'present figure'):
+        fig.add_subplot(ax)
 
 
 @image_comparison(['figure_suptitle'])
@@ -426,6 +446,12 @@ def test_invalid_figure_add_axes():
     with pytest.raises(TypeError, match="multiple values for argument 'rect'"):
         fig.add_axes([0, 0, 1, 1], rect=[0, 0, 1, 1])
 
+    _, ax = plt.subplots()
+    with pytest.raises(ValueError,
+                       match="The Axes must have been created in the present "
+                             "figure"):
+        fig.add_axes(ax)
+
 
 def test_subplots_shareax_loglabels():
     fig, axs = plt.subplots(2, 2, sharex=True, sharey=True, squeeze=False)
