Add support for categorical colorbars

dstansby · dstansby · commit 979b099ed59f · 2024-03-10T10:24:50.000Z
diff --git a/doc/api/next_api_changes/behavior/27721-DS.rst b/doc/api/next_api_changes/behavior/27721-DS.rst
@@ -14,3 +14,7 @@ convert data of more than one dimension (e.g. when plotting images the data is 2
 If a conversion interface raises an error when given ``None`` or 2D data as described
 above, this error will be re-raised when a user tries to use one of the newly supported
 plotting methods with unit-ful data.
+
+If you have a custom conversion interface you want to forbid using with image data, the
+`~.units.ConversionInterface` methods that accept a ``units`` parameter should raise
+a `matplotlib.units.ConversionError` when given ``units=None``.
diff --git a/lib/matplotlib/category.py b/lib/matplotlib/category.py
@@ -101,6 +101,8 @@ def default_units(data, axis):
             object storing string to integer mapping
         """
         # the conversion call stack is default_units -> axis_info -> convert
+        if axis is None:
+            return UnitData(data)
         if axis.units is None:
             axis.set_units(UnitData(data))
         else:
@@ -208,7 +210,7 @@ def update(self, data):
         TypeError
             If elements in *data* are neither str nor bytes.
         """
-        data = np.atleast_1d(np.array(data, dtype=object))
+        data = np.atleast_1d(np.array(data, dtype=object).ravel())
         # check if convertible to number:
         convertible = True
         for val in OrderedDict.fromkeys(data):
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -408,6 +408,9 @@ def _strip_units(self, A):
         try:
             self._units = self._converter.default_units(A, None)
         except Exception as e:
+            if isinstance(e, munits.ConversionError):
+                raise e
+
             raise RuntimeError(
                 f'{self._converter} failed when trying to return the default units for '
                 'this image. This may be because support has not been '
@@ -417,6 +420,9 @@ def _strip_units(self, A):
         try:
             A = self._converter.convert(A, self._units, None)
         except Exception as e:
+            if isinstance(e, munits.ConversionError):
+                raise e
+
             raise munits.ConversionError(
                 f'{self._converter} failed when trying to convert the units for this '
                 'image. This may be because support has not been implemented '
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -333,9 +333,10 @@ def test_mappable_units(quantity_converter):
     mappable = ax.pcolor(x, y, data, vmin=vmin, vmax=vmax)
     fig.colorbar(mappable, ax=ax, extend="min")
 
-    # pcolormesh + horizontal colorbar
+    # pcolormesh + horizontal colorbar + categorical
+    data = [["one", "two"], ["three", "four"]]
     ax = axs[1, 0]
-    mappable = ax.pcolormesh(x, y, data, vmin=vmin, vmax=vmax)
+    mappable = ax.pcolormesh(x, y, data)
     fig.colorbar(mappable, ax=ax, orientation="horizontal", extend="min")
 
     axs[1, 1].axis("off")
