FIX: make manual axes more like normal

jklymak · jklymak · commit 4999aa754ca3 · 2021-05-02T12:29:37.000-07:00
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -38,6 +38,7 @@
 from matplotlib import _api, collections, cm, colors, contour, ticker
 from matplotlib.axes._base import _TransformedBoundsLocator
 from matplotlib.axes._axes import Axes
+from matplotlib.axes._subplots import SubplotBase
 import matplotlib.artist as martist
 import matplotlib.patches as mpatches
 import matplotlib.path as mpath
@@ -235,24 +236,44 @@ class ColorbarAxes(Axes):
     Users should not normally instantiate this class, but it is the class
     returned by ``cbar = fig.colorbar(im); cax = cbar.ax``.
     """
-    def __init__(self, parent):
+    def __init__(self, parent, userax=False):
         """
+        Parameters
+        ----------
+        parent : Axes
+            Axes that specifies the position of the colorbar.
+        userax : boolean
+            True if the user passed `.Figure.colorbar` the axes manually.
         """
-        inner_ax = parent.inset_axes([0, 0, 1, 1])
-        #self.__class__ = type(inner_ax.__class__.__name__+"ColorBar",
-        #                      (self.__class__, inner_ax.__class__),
-        #                      {})
+        
+        if userax:
+            # copy position:
+            fig = parent.figure
+            parent_ax = fig.add_axes(parent.get_position())
+            # copy the locator if one exists:
+            parent_ax._axes_locator = parent._axes_locator
+            # remove the parent from the figure...
+            parent.remove()    
+        else:
+            parent_ax = parent
+
+        inner_ax = parent_ax.inset_axes([0, 0, 1, 1])
         self.__dict__ = inner_ax.__dict__
-        self.parent_ax = parent
+        
+        self.parent_ax = parent_ax
         self.inner_ax = inner_ax
         self.parent_ax.xaxis.set_visible(False)
         self.parent_ax.yaxis.set_visible(False)
         self.parent_ax.set_facecolor('none')
-        # map some features to the parent so users have access...
         self.parent_ax.tick_params = self.inner_ax.tick_params
-
+        self.parent_ax.set_xticks = self.inner_ax.set_xticks
+        self.parent_ax.set_yticks = self.inner_ax.set_yticks
         for attr in ["get_position", "set_position", "set_aspect"]:
             setattr(self, attr, getattr(self.parent_ax, attr))
+        if userax:
+            # point the parent's methods all at this axes...
+            parent.__dict__ = self.__dict__
+
 
     def _set_inner_bounds(self, bounds):
         """
@@ -379,6 +400,7 @@ def __init__(self, ax, cmap=None,
                  extendfrac=None,
                  extendrect=False,
                  label='',
+                 userax=False,
                  ):
         _api.check_isinstance([colors.Colormap, None], cmap=cmap)
         _api.check_in_list(
@@ -390,10 +412,8 @@ def __init__(self, ax, cmap=None,
             ['uniform', 'proportional'], spacing=spacing)
 
         # wrap the axes so that it can be positioned as an inset axes:
-        ax = ColorbarAxes(ax)
+        ax = ColorbarAxes(ax, userax=userax)
         self.ax = ax
-        # Bind some methods to the axes to warn users against using them.
-        ax.set_xticks = ax.set_yticks = _set_ticks_on_axis_warn
         ax.set(navigate=False)
 
         if cmap is None:
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -1156,21 +1156,28 @@ def colorbar(self, mappable, cax=None, ax=None, use_gridspec=True, **kw):
                     "to colorbar().")
 
         # Store the value of gca so that we can set it back later on.
-        current_ax = self.gca()
         if cax is None:
+            current_ax = self.gca()
+            userax = False
             if (use_gridspec and isinstance(ax, SubplotBase)
                     and not self.get_constrained_layout()):
                 cax, kw = cbar.make_axes_gridspec(ax, **kw)
             else:
                 cax, kw = cbar.make_axes(ax, **kw)
+        else:
+            userax = True
 
         # need to remove kws that cannot be passed to Colorbar
         NON_COLORBAR_KEYS = ['fraction', 'pad', 'shrink', 'aspect', 'anchor',
                              'panchor']
         cb_kw = {k: v for k, v in kw.items() if k not in NON_COLORBAR_KEYS}
+        if userax:
+            cb_kw['userax'] = True
+
         cb = cbar.Colorbar(cax, mappable, **cb_kw)
 
-        self.sca(current_ax)
+        if not userax:
+            self.sca(current_ax)
         self.stale = True
         return cb
 
