Deal with axes_grid differences

jklymak · jklymak · commit fb768c7b90b9 · 2021-06-22T14:23:07.000-07:00
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -239,15 +239,19 @@ def __init__(self, parent, userax=True):
                 parent._axes.add_child_axes(outer_ax)
                 outer_ax._axes.child_axes.remove(parent)
             else:
-                parent.remove()
+                try:
+                    parent.remove()
+                except ValueError:
+                    pass  # Already removed
         else:
             outer_ax = parent
 
-        # swap axes in the stack:
-        fig._localaxes.remove(outer_ax)
-        fig._axstack.remove(outer_ax)
-        fig._localaxes.add(self)
-        fig._axstack.add(self)
+        # swap axes in the stack if its in there:
+        if outer_ax in fig._localaxes:
+            fig._localaxes.remove(outer_ax)
+            fig._axstack.remove(outer_ax)
+            fig._localaxes.add(self)
+            fig._axstack.add(self)
         inner_ax = outer_ax.inset_axes([0, 0, 1, 1])
         self.__dict__.update(inner_ax.__dict__)
 
@@ -261,8 +265,8 @@ def __init__(self, parent, userax=True):
         self.outer_ax.set_xticks = self.inner_ax.set_xticks
         self.outer_ax.set_yticks = self.inner_ax.set_yticks
         for attr in ["get_position", "set_aspect",
-                    "_remove_method", "_set_position",
-                     "set_position"]:
+                     "_remove_method", "_set_position",
+                     "set_position", "cla", "draw"]:
             setattr(self, attr, getattr(self.outer_ax, attr))
         self._colorbar_info = None  # used for mpl-created axes
         if hasattr(self.outer_ax, "get_subplotspec"):
@@ -272,7 +276,11 @@ def __init__(self, parent, userax=True):
         if userax:
             self._colorbar_info = 'user'
             # point the parent's methods all at this axes...
+            origdict = parent.__dict__
             parent.__dict__ = self.__dict__
+            for key in origdict.keys():
+                if key not in parent.__dict__:
+                    parent.__dict__[key] = origdict[key]
 
     def _set_inner_bounds(self, bounds):
         """
@@ -281,9 +289,6 @@ def _set_inner_bounds(self, bounds):
         self.inner_ax._axes_locator = _TransformedBoundsLocator(
             bounds, self.outer_ax.transAxes)
 
-    def draw(self, renderer):
-        self.outer_ax.draw(renderer)
-
 
 class _ColorbarSpine(mspines.Spine):
     def __init__(self, axes):
diff --git a/lib/mpl_toolkits/axes_grid1/axes_grid.py b/lib/mpl_toolkits/axes_grid1/axes_grid.py
@@ -3,7 +3,6 @@
 
 import numpy as np
 
-import matplotlib as mpl
 from matplotlib import _api
 from matplotlib.gridspec import SubplotSpec
 
@@ -29,28 +28,19 @@ def colorbar(self, mappable, *, ticks=None, **kwargs):
         orientation = (
             "horizontal" if self.orientation in ["top", "bottom"] else
             "vertical")
-        kwargs['userax'] = False
-        cb = mpl.colorbar.Colorbar(
-            self, mappable, orientation=orientation, ticks=ticks, **kwargs)
-        self._config_axes()
+        cb = self.figure.colorbar(mappable, cax=self, orientation=orientation,
+                                  ticks=ticks, **kwargs)
         return cb
 
-    def _config_axes(self):
-        """Make an axes patch and outline."""
-        ax = self
-        ax.set_navigate(False)
-        ax.axis[:].toggle(all=False)
-        b = self._default_label_on
-        ax.axis[self.orientation].toggle(all=b)
-
     def toggle_label(self, b):
         self._default_label_on = b
         axis = self.axis[self.orientation]
         axis.toggle(ticklabels=b, label=b)
 
     def cla(self):
+        orientation = self.orientation
         super().cla()
-        self._config_axes()
+        self.orientation = orientation
 
 
 class CbarAxes(CbarAxesBase, Axes):
