Add location keyword argument to Colorbar

oscargus · oscargus · commit ee796c5e4ed3 · 2022-11-02T14:31:51.000+01:00
diff --git a/doc/users/next_whats_new/colorbar_has_location_argument.rst b/doc/users/next_whats_new/colorbar_has_location_argument.rst
@@ -0,0 +1,25 @@
+``colorbar`` now has a *location* keyword argument
+==================================================
+
+The ``colorbar`` method now supports a *location* keyword argument to more
+easily position the color bar. This is useful when providing your own inset
+axes using the *cax* keyword argument and behaves similar to the case where
+axes are not provided (where the *location* keyword is passed through).
+*orientation* and *ticklocation* are no longer required as they are
+determined by *location*. *ticklocation* can still be provided if the
+automatic setting is not preferred. (*orientation* can also be provided but
+must be compatible with the *location*.)
+
+An example is:
+
+.. plot::
+    :include-source: true
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+    rng = np.random.default_rng(19680801)
+    imdata = rng.random((10, 10))
+    fig, ax = plt.subplots()
+    im = ax.imshow(imdata)
+    fig.colorbar(im, cax=ax.inset_axes([0, 1.05, 1, 0.05]),
+                 location='top')
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -42,6 +42,11 @@
     of the colorbar, as that also determines the *orientation*; passing
     incompatible values for *location* and *orientation* raises an exception.
 
+ticklocation : {'auto', 'left', 'right', 'top', 'bottom'}
+    The location of the colorbar ticks. The *ticklocation* location must match
+    *orientation*. For example, a horizontal colorbar can only have ticks at
+    the top or the bottom.
+
 fraction : float, default: 0.15
     Fraction of original axes to use for colorbar.
 
@@ -246,13 +251,18 @@ class Colorbar:
     alpha : float
         The colorbar transparency between 0 (transparent) and 1 (opaque).
 
-    orientation : {'vertical', 'horizontal'}
+    orientation : None or {'vertical', 'horizontal'}
+        If None, use the value determined by *location*. If both
+        *orientation* and *location* are None then defaults to 'vertical'.
 
     ticklocation : {'auto', 'left', 'right', 'top', 'bottom'}
 
     drawedges : bool
 
     filled : bool
+
+    location : None or {'left', 'right', 'top', 'bottom'}
+
     %(_colormap_kw_doc)s
     """
 
@@ -264,7 +274,7 @@ def __init__(self, ax, mappable=None, *, cmap=None,
                  alpha=None,
                  values=None,
                  boundaries=None,
-                 orientation='vertical',
+                 orientation=None,
                  ticklocation='auto',
                  extend=None,
                  spacing='uniform',  # uniform or proportional
@@ -275,6 +285,7 @@ def __init__(self, ax, mappable=None, *, cmap=None,
                  extendfrac=None,
                  extendrect=False,
                  label='',
+                 location=None,
                  ):
 
         if mappable is None:
@@ -305,14 +316,23 @@ def __init__(self, ax, mappable=None, *, cmap=None,
         mappable.colorbar_cid = mappable.callbacks.connect(
             'changed', self.update_normal)
 
+        location_orientation = _get_orientation_from_location(location)
+
         _api.check_in_list(
-            ['vertical', 'horizontal'], orientation=orientation)
+            [None, 'vertical', 'horizontal'], orientation=orientation)
         _api.check_in_list(
             ['auto', 'left', 'right', 'top', 'bottom'],
             ticklocation=ticklocation)
         _api.check_in_list(
             ['uniform', 'proportional'], spacing=spacing)
 
+        if location_orientation is not None and orientation is not None:
+            if location_orientation != orientation:
+                raise TypeError(
+                    "location and orientation are mutually exclusive")
+        else:
+            orientation = orientation or location_orientation or "vertical"
+
         self.ax = ax
         self.ax._axes_locator = _ColorbarAxesLocator(self)
 
@@ -365,7 +385,8 @@ def __init__(self, ax, mappable=None, *, cmap=None,
         self.__scale = None  # linear, log10 for now.  Hopefully more?
 
         if ticklocation == 'auto':
-            ticklocation = 'bottom' if orientation == 'horizontal' else 'right'
+            ticklocation = _get_ticklocation_from_orientation(
+                orientation) if location is None else location
         self.ticklocation = ticklocation
 
         self.set_label(label)
@@ -1330,25 +1351,36 @@ def drag_pan(self, button, key, x, y):
 
 def _normalize_location_orientation(location, orientation):
     if location is None:
-        location = _api.check_getitem(
-            {None: "right", "vertical": "right", "horizontal": "bottom"},
-            orientation=orientation)
+        location = _get_ticklocation_from_orientation(orientation)
     loc_settings = _api.check_getitem({
-        "left":   {"location": "left", "orientation": "vertical",
-                   "anchor": (1.0, 0.5), "panchor": (0.0, 0.5), "pad": 0.10},
-        "right":  {"location": "right", "orientation": "vertical",
-                   "anchor": (0.0, 0.5), "panchor": (1.0, 0.5), "pad": 0.05},
-        "top":    {"location": "top", "orientation": "horizontal",
-                   "anchor": (0.5, 0.0), "panchor": (0.5, 1.0), "pad": 0.05},
-        "bottom": {"location": "bottom", "orientation": "horizontal",
-                   "anchor": (0.5, 1.0), "panchor": (0.5, 0.0), "pad": 0.15},
+        "left":   {"location": "left", "anchor": (1.0, 0.5),
+                   "panchor": (0.0, 0.5), "pad": 0.10},
+        "right":  {"location": "right", "anchor": (0.0, 0.5),
+                   "panchor": (1.0, 0.5), "pad": 0.05},
+        "top":    {"location": "top", "anchor": (0.5, 0.0),
+                   "panchor": (0.5, 1.0), "pad": 0.05},
+        "bottom": {"location": "bottom", "anchor": (0.5, 1.0),
+                   "panchor": (0.5, 0.0), "pad": 0.15},
     }, location=location)
+    loc_settings["orientation"] = _get_orientation_from_location(location)
     if orientation is not None and orientation != loc_settings["orientation"]:
         # Allow the user to pass both if they are consistent.
         raise TypeError("location and orientation are mutually exclusive")
     return loc_settings
 
 
+def _get_orientation_from_location(location):
+    return _api.check_getitem(
+        {None: None, "left": "vertical", "right": "vertical",
+         "top": "horizontal", "bottom": "horizontal"}, location=location)
+
+
+def _get_ticklocation_from_orientation(orientation):
+    return _api.check_getitem(
+        {None: "right", "vertical": "right", "horizontal": "bottom"},
+        orientation=orientation)
+
+
 @_docstring.interpd
 def make_axes(parents, location=None, orientation=None, fraction=0.15,
               shrink=1.0, aspect=20, **kwargs):
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -1149,3 +1149,34 @@ def test_title_text_loc():
     # colorbar axes, including its extend triangles....
     assert (cb.ax.title.get_window_extent(fig.canvas.get_renderer()).ymax >
             cb.ax.spines['outline'].get_window_extent().ymax)
+
+
+@check_figures_equal(extensions=["png"])
+def test_passing_location(fig_ref, fig_test):
+    ax_ref = fig_ref.add_subplot()
+    im = ax_ref.imshow([[0, 1], [2, 3]])
+    ax_ref.figure.colorbar(im, cax=ax_ref.inset_axes([0, 1.05, 1, 0.05]),
+                           orientation="horizontal", ticklocation="top")
+    ax_test = fig_test.add_subplot()
+    im = ax_test.imshow([[0, 1], [2, 3]])
+    ax_test.figure.colorbar(im, cax=ax_test.inset_axes([0, 1.05, 1, 0.05]),
+                            location="top")
+
+
+@pytest.mark.parametrize("kwargs,error,message", [
+    ({'location': 'top', 'orientation': 'vertical'}, TypeError,
+     "location and orientation are mutually exclusive"),
+    ({'location': 'top', 'orientation': 'vertical', 'cax': True}, TypeError,
+     "location and orientation are mutually exclusive"),  # Different to above
+    ({'ticklocation': 'top', 'orientation': 'vertical', 'cax': True},
+     ValueError, "'top' is not a valid value for position"),
+    ({'location': 'top', 'extendfrac': (0, None)}, ValueError,
+     "invalid value for extendfrac"),
+    ])
+def test_colorbar_errors(kwargs, error, message):
+    fig, ax = plt.subplots()
+    im = ax.imshow([[0, 1], [2, 3]])
+    if kwargs.get('cax', None) is True:
+        kwargs['cax'] = ax.inset_axes([0, 1.05, 1, 0.05])
+    with pytest.raises(error, match=message):
+        fig.colorbar(im, **kwargs)
