TST: and a few fixes

jklymak · jklymak · commit 439a0256cfc3 · 2021-05-06T17:37:50.000-07:00
diff --git a/doc/api/next_api_changes/behavior/20054-JMK.rst b/doc/api/next_api_changes/behavior/20054-JMK.rst
@@ -6,16 +6,9 @@ The axes used to place a colorbar is now wrapped by a new parent class
 
     cb = fig.colorbar(im, cax=cax)
 
-Now ``cb.ax`` returns the parent ``ColorbarAxes``, and to get
-``cax`` back, you can do::
-
-    cb.ax.parent_axes
-
-The parent axes (``cb.ax.parent_axes``) handles the
-placement of the colorbar axes object, and its aspect ratio, but the child
-axes (``cb.ax``)  handles all the tick parameters, labelling, and holds the
-artists that make up the colorbar.  We have attempted to dispatch the
-appropriate colorbar methods to the appropriate axes.  
+This means that ``cb.ax`` is no longer the same object as ``cax``. However, 
+we map all the methods from ``cb.ax`` onto ``cax`` so ``cax`` should remain 
+functionally the same as ``cb.ax``.  
 
 Colorbar lines no longer clipped
 ================================
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -377,6 +377,9 @@ class ColorbarBase:
     extendrec
 
     label : str
+
+    userax : boolean
+        Whether the user created the axes or not.  Default True
     """
 
     n_rasterize = 50  # rasterize solids if number of colors >= n_rasterize
@@ -398,7 +401,7 @@ def __init__(self, ax, cmap=None,
                  extendfrac=None,
                  extendrect=False,
                  label='',
-                 userax=False,
+                 userax=True,
                  ):
         _api.check_isinstance([colors.Colormap, None], cmap=cmap)
         _api.check_in_list(
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -1158,25 +1158,23 @@ def colorbar(self, mappable, cax=None, ax=None, use_gridspec=True, **kw):
         # Store the value of gca so that we can set it back later on.
         if cax is None:
             current_ax = self.gca()
-            userax = False
+            kw['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
+            kw['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)
 
-        if not userax:
+        if not kw['userax']:
             self.sca(current_ax)
         self.stale = True
         return cb
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -12,6 +12,8 @@
 from matplotlib.colorbar import ColorbarBase
 from matplotlib.ticker import FixedLocator
 
+from matplotlib.testing.decorators import check_figures_equal
+
 
 def _get_cmap_norms():
     """
@@ -724,3 +726,20 @@ def test_colorbar_change_lim_scale():
     pc = ax[1].pcolormesh(np.arange(100).reshape(10, 10)+1)
     cb = fig.colorbar(pc, ax=ax[1], extend='both')
     cb.ax.set_ylim([20, 90])
+
+
+@check_figures_equal(extensions=["png"])
+def test_axes_handles_same_functions(fig_ref, fig_test):
+    # prove that cax and cb.ax are functionally the same
+    for nn, fig in enumerate([fig_ref, fig_test]):
+        ax = fig.add_subplot()
+        pc = ax.pcolormesh(np.ones(300).reshape(10, 30))
+        cax = fig.add_axes([0.9, 0.1, 0.03, 0.8])
+        cb = fig.colorbar(pc, cax=cax)
+        if nn == 0:
+            caxx = cax
+        else:
+            caxx = cb.ax
+        caxx.set_yticks(np.arange(20))
+        caxx.set_yscale('log')
+        caxx.set_position([0.92, 0.1, 0.02, 0.7])
