ENH: Add Axes.colorbar() method

Author: timhoffm

doc/api/axes_api.rst

@@ -314,8 +314,8 @@ Axis limits and direction
    Axes.set_ybound
    Axes.get_ybound
 
-Axis labels, title, and legend
-------------------------------
+Axis labels and title
+---------------------
 
 .. autosummary::
    :toctree: _as_gen
@@ -330,9 +330,6 @@ Axis labels, title, and legend
 
    Axes.set_title
    Axes.get_title
-   Axes.legend
-   Axes.get_legend
-   Axes.get_legend_handles_labels
 
 Axis scales
 -----------
@@ -438,6 +435,20 @@ Ticks and tick labels
    Axes.locator_params
 
 
+Legend and colorbar
+===================
+
+.. autosummary::
+   :toctree: _as_gen
+   :template: autosummary.rst
+   :nosignatures:
+
+   Axes.legend
+   Axes.get_legend
+   Axes.get_legend_handles_labels
+   Axes.colorbar
+
+
 Units
 =====
 

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 if you want to place a colorbar
+        right next to an Axes. ``ax.colorbar(mappable)`` is equivalent to
+        ``fig.colorbar(mappable, ax=ax)``. In particular, if you have exactly one
+        mappable in the Axes, the general
+
+        ::
+            im = ax.imshow(data)
+            fig.colorbar(im, ax=ax)
+
+        can be written more concisely as ::
+
+            ax.imshow(data)
+            ax.colorbar()
+
+        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/axes/_axes.pyi

@@ -11,7 +11,9 @@ from matplotlib.collections import (
     EventCollection,
     QuadMesh,
 )
-from matplotlib.colorizer import Colorizer
+from matplotlib.cm import ScalarMappable
+from matplotlib.colorbar import Colorbar
+from matplotlib.colorizer import Colorizer, ColorizingArtist
 from matplotlib.colors import Colormap, Normalize
 from matplotlib.container import (
     BarContainer, PieContainer, ErrorbarContainer, StemContainer)
@@ -77,6 +79,13 @@ class Axes(_AxesBase):
     @overload
     def legend(self, *, loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
 
+    def colorbar(
+        self,
+        mappable: ScalarMappable | ColorizingArtist | None = ...,
+        *,
+        use_gridspec: bool = ...,
+        **kwargs
+    ) -> Colorbar: ...
     def inset_axes(
         self,
         bounds: tuple[float, float, float, float],

lib/matplotlib/tests/test_axes.py

@@ -9329,6 +9329,27 @@ def test_automatic_legend():
     assert ax.get_yticklabels()[0].get_text() == 'b'
 
 
+def test_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
+
+    # in case of multiple mappables: bare colorbar() fails, but passing a mappable works
+    im = 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()
+    cb3 = ax3.colorbar(im)
+    assert cb3.mappable is im
+
+
 def test_plot_errors():
     with pytest.raises(TypeError, match=r"plot\(\) got an unexpected keyword"):
         plt.plot([1, 2, 3], x=1)