ENH: add colorbar method to axes

jklymak · jklymak · commit 008410196108 · 2018-09-28T21:01:07.000-07:00
diff --git a/doc/users/next_whats_new/2018-09-28-JMK.rst b/doc/users/next_whats_new/2018-09-28-JMK.rst
@@ -0,0 +1,8 @@
+:orphan:
+
+New `.axes.Axes.colorbar` method for attaching colorbar to single axes
+``````````````````````````````````````````````````````````````````````
+
+A simple wrapper for `.Figure.colorbar` that allows the user to omit the
+``ax`` kwarg. The new ``ax.colorbar(mappable)`` is equivalent to
+``fig.colorbar(mappable, ax=ax)``.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -418,6 +418,21 @@ def legend(self, *args, **kwargs):
     def _remove_legend(self, legend):
         self.legend_ = None
 
+    @docstring.dedent_interpd
+    def colorbar(self, mappable, cax=None, use_gridspec=True, **kw):
+        """
+        Create a colorbar for a ScalarMappable instance, *mappable*, with the
+        axes as the parent.
+
+        Documentation for the pyplot thin wrapper:
+        %(colorbar_doc)s
+        """
+        ax = kw.pop('ax', None)
+        if ax is not None:
+            warnings.warn('Supplied ax argument ignored')
+        self.figure.colorbar(mappable,  ax=self, cax=cax,
+                             use_gridspec=use_gridspec, **kw)
+
     def inset_axes(self, bounds, *, transform=None, zorder=5,
             **kwargs):
         """
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -396,3 +396,19 @@ def test_colorbar_axes_kw():
     plt.imshow(([[1, 2], [3, 4]]))
     plt.colorbar(orientation='horizontal', fraction=0.2, pad=0.2, shrink=0.5,
                  aspect=10, anchor=(0., 0.), panchor=(0., 1.))
+
+
+def test_colorbar_axes():
+    fig, ax = plt.subplots()
+    pcm = ax.pcolormesh(np.random.random((32, 32)))
+    # smoketest that this works:
+    ax.colorbar(pcm)
+    with pytest.warns(UserWarning) as record:
+        ax.colorbar(pcm, ax=ax)
+        assert len(record) == 1
+
+    fig, ax = plt.subplots()
+    pcm = ax.pcolormesh(np.random.random((32, 32)))
+    # smoketest that this works:
+    ax.colorbar(pcm, orientation='horizontal', fraction=0.2, pad=0.2, shrink=0.5,
+                 aspect=10, anchor=(0., 0.), panchor=(0., 1.))
diff --git a/lib/mpl_toolkits/axes_grid1/colorbar.py b/lib/mpl_toolkits/axes_grid1/colorbar.py
@@ -109,7 +109,8 @@
     the :class:`~matplotlib.image.Image`,
     :class:`~matplotlib.contour.ContourSet`, etc. to
     which the colorbar applies; this argument is mandatory for the
-    :meth:`~matplotlib.figure.Figure.colorbar` method but optional for the
+    :meth:`~matplotlib.figure.Figure.colorbar`
+    and :meth:`~matplotlib.axes.Axes.colorbar` methods but optional for the
     :func:`~matplotlib.pyplot.colorbar` function, which sets the
     default to the current image.
 
@@ -119,7 +120,7 @@
     None | axes object into which the colorbar will be drawn
   *ax*
     None | parent axes object from which space for a new
-    colorbar axes will be stolen
+    colorbar axes will be stolen (ignored for `.axes.Axes.colorbar`)
 
 
 Additional keyword arguments are of two kinds:
