Allow ticks argument to take units

dstansby · dstansby · commit 7234bf59b495 · 2024-01-30T21:32:16.000Z
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -370,7 +370,6 @@ def __init__(self, ax, mappable=None, *, cmap=None,
         self.solids_patches = []
         self.lines = []
 
-
         for spine in self.ax.spines.values():
             spine.set_visible(False)
         self.outline = self.ax.spines['outline'] = _ColorbarSpine(self.ax)
@@ -397,9 +396,12 @@ def __init__(self, ax, mappable=None, *, cmap=None,
         self._reset_locator_formatter_scale()
         self._set_units_from_mappable()
 
+        if ticks is not None and self._converter is not None:
+            ticks = self._converter.convert(ticks, self._units, self._long_axis())
+
         if np.iterable(ticks):
             self._locator = ticker.FixedLocator(ticks, nbins=len(ticks))
-        else:
+        elif isinstance(ticks, ticker.Locator):
             self._locator = ticks
 
         if isinstance(format, str):
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -305,14 +305,18 @@ def test_mappable_units(quantity_converter):
     cbar = fig.colorbar(mappable, ax=ax, extend="min")
 
     # pcolor
+    #
     # Use datetime to check that the locator/formatter is set correctly
+    # Also test ticks argument to colorbar
     data = [[datetime(2024, 1, 1), datetime(2024, 1, 2)],
             [datetime(2024, 1, 3), datetime(2024, 1, 4)]]
     vmin = datetime(2024, 1, 2)
     vmax = datetime(2024, 1, 5)
     ax = axs[0, 1]
     mappable = ax.pcolor(x, y, data, vmin=vmin, vmax=vmax)
-    fig.colorbar(mappable, ax=ax, extend="min")
+
+    ticks = [datetime(2024, 1, 2), datetime(2024, 1, 3), datetime(2024, 1, 5)]
+    fig.colorbar(mappable, ax=ax, extend="min", ticks=ticks)
 
     # pcolormesh + horizontal colorbar + categorical
     data = [["one", "two"], ["three", "four"]]
