Set units on colorbar axes

dstansby · dstansby · commit f20a66a21e9c · 2024-02-20T11:44:54.000Z
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -291,7 +291,7 @@ def __init__(self, ax, mappable=None, *, cmap=None,
                  drawedges=False,
                  extendfrac=None,
                  extendrect=False,
-                 label='',
+                 label=None,
                  location=None,
                  ):
 
@@ -370,6 +370,7 @@ 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)
@@ -391,8 +392,10 @@ def __init__(self, ax, mappable=None, *, cmap=None,
                 orientation) if location is None else location
         self.ticklocation = ticklocation
 
-        self.set_label(label)
+        if label is not None:
+            self.set_label(label)
         self._reset_locator_formatter_scale()
+        self._set_units_from_mappable()
 
         if np.iterable(ticks):
             self._locator = ticker.FixedLocator(ticks, nbins=len(ticks))
@@ -1331,6 +1334,27 @@ def drag_pan(self, button, key, x, y):
             elif self.orientation == 'vertical':
                 self.norm.vmin, self.norm.vmax = points[:, 1]
 
+    def _set_units_from_mappable(self):
+        """
+        Set the colorbar locator and formatter if the mappable has units.
+        """
+        self._units = self.mappable._units
+        self._converter = self.mappable._converter
+        if self._converter is not None:
+
+            axis = self._long_axis()
+            info = self._converter.axisinfo(self._units, axis)
+
+            if info is not None:
+                if info.majloc is not None:
+                    self.locator = info.majloc
+                if info.minloc is not None:
+                    self.minorlocator = info.minloc
+                if info.majfmt is not None:
+                    self.formatter = info.majfmt
+                if info.minfmt is not None:
+                    self.minorformatter = info.minfmt
+
 
 ColorbarBase = Colorbar  # Backcompat API
 
diff --git a/lib/matplotlib/colorbar.pyi b/lib/matplotlib/colorbar.pyi
@@ -59,7 +59,7 @@ class Colorbar:
         drawedges: bool = ...,
         extendfrac: Literal["auto"] | float | Sequence[float] | None = ...,
         extendrect: bool = ...,
-        label: str = ...,
+        label: str | None = ...,
         location: Literal["left", "right", "top", "bottom"] | None = ...
     ) -> None: ...
     @property
