Add location keyword argument to Colorbar

oscargus · oscargus · commit f2cd4da51fb3 · 2022-06-14T11:25:24.000+02:00
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 value 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.
 
@@ -321,13 +326,17 @@ 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, 'vertical'.
 
     ticklocation : {'auto', 'left', 'right', 'top', 'bottom'}
 
     drawedges : bool
 
     filled : bool
+
+    location : None or {'left', 'right', 'top', 'bottom'}
     %s
     """
 
@@ -339,7 +348,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
@@ -350,6 +359,7 @@ def __init__(self, ax, mappable=None, *, cmap=None,
                  extendfrac=None,
                  extendrect=False,
                  label='',
+                 location=None,
                  ):
 
         if mappable is None:
@@ -380,14 +390,26 @@ def __init__(self, ax, mappable=None, *, cmap=None,
         mappable.colorbar_cid = mappable.callbacks.connect(
             'changed', self.update_normal)
 
+        location_orientation = _api.check_getitem(
+            {None: None, "left": "vertical", "right": "vertical",
+             "top": "horizontal", "bottom": "horizontal"},
+            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)
 
@@ -445,7 +467,11 @@ 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'
+            if location is None:
+                ticklocation = (
+                    'bottom' if orientation == 'horizontal' else 'right')
+            else:
+                ticklocation = location
         self.ticklocation = ticklocation
 
         self.set_label(label)
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -1049,3 +1049,15 @@ 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")
