ENH: Add Axes.colorbar() method

Author: timhoffm

lib/matplotlib/axes/_axes.py

@@ -355,6 +355,73 @@ def legend(self, *args, **kwargs):
     def _remove_legend(self, legend):
         self.legend_ = None
 
+    def colorbar(self, mappable=None, *, use_gridspec=True, **kwargs):
+        """
+        Add a colorbar next to the Axes.
+
+        Parameters
+        ----------
+        mappable : `matplotlib.colorizer.ColorizingArtist`, optional
+            The `matplotlib.colorizer.ColorizingArtist` (i.e., `.AxesImage`,
+            `.ContourSet`, etc.) described by this colorbar.
+
+            If not given, the mappable is inferred from the Axes. If there is exactly
+            one image or collection, this is used as mappable. Otherwise, an error is
+            raised and you must specify the mappable explicitly.
+
+            Note that one can create a `.colorizer.ColorizingArtist` "on-the-fly"
+            to generate colorbars not attached to a previously drawn artist, e.g.
+            ::
+
+                cr = colorizer.Colorizer(norm=norm, cmap=cmap)
+                ax.colorbar(colorizer.ColorizingArtist(cr))
+
+        Returns
+        -------
+        colorbar : `~matplotlib.colorbar.Colorbar`
+
+        Other Parameters
+        ----------------
+        use_gridspec : bool, optional
+            If *ax* is positioned with a subplotspec and *use_gridspec*
+            is ``True``, then *cax* is also positioned with a subplotspec.
+
+        %(_make_axes_kw_doc)s
+        %(_colormap_kw_doc)s
+
+        Notes
+        -----
+        This method is a convenience shortcut for the most common case that you have
+        one mappable in the Axes and want to create a colorbar for it, which is
+        located right beside the Axes.
+
+        ::
+            ax.imshow(data)
+            ax.colorbar()
+
+        is equivalent to
+
+        ::
+            im = ax.imshow(data)
+            fig.colorbar(im, ax=ax)
+
+        Use `.Figure.colorbar` if you need more control on placing the colorbar.
+        """
+        if mappable is None:
+            # autodetect the mappable
+            colormapped_artists = [*self.images, *self.collections]
+            if len(colormapped_artists) != 1:
+                raise RuntimeError(
+                    "Axes.colormap() requires exactly one colormapped Artist in the "
+                    f"Axes, but found {len(colormapped_artists)}. In ambiguous cases, "
+                    "please specify the mappable for the colorbar explicitly."
+                )
+            mappable = colormapped_artists[0]
+
+        return self.figure.colorbar(
+            mappable, ax=self, use_gridspec=use_gridspec, **kwargs
+        )
+
     def inset_axes(self, bounds, *, transform=None, zorder=5, **kwargs):
         """
         Add a child inset Axes to this existing Axes.

lib/matplotlib/tests/test_axes.py

@@ -9329,6 +9329,24 @@ def test_automatic_legend():
     assert ax.get_yticklabels()[0].get_text() == 'b'
 
 
+def test_automatic_colorbar():
+    """Test that Axes.colorbar() without arguments uses the only mappable."""
+    fig, (ax1, ax2, ax3) = plt.subplots(3, 1)
+
+    im = ax1.imshow([[0, 1], [2, 3]])
+    cb = ax1.colorbar()
+    assert cb.mappable is im
+
+    path_collection = ax2.scatter([0, 1], [0, 1], c=[0, 1])
+    cb2 = ax2.colorbar()
+    assert cb2.mappable is path_collection
+
+    ax3.imshow([[0, 1], [2, 3]])
+    ax3.scatter([0, 1], [0, 1], c=[0, 1])
+    with pytest.raises(RuntimeError, match="requires exactly one colormapped Artist"):
+        ax3.colorbar()
+
+
 def test_plot_errors():
     with pytest.raises(TypeError, match=r"plot\(\) got an unexpected keyword"):
         plt.plot([1, 2, 3], x=1)