Add support for categorical colorbars

dstansby · dstansby · commit f368afddcd5b · 2024-01-30T21:32:16.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
@@ -529,6 +529,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 '
@@ -538,6 +541,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_category.py b/lib/matplotlib/tests/test_category.py
@@ -9,6 +9,7 @@
 import matplotlib.pyplot as plt
 import matplotlib.category as cat
 from matplotlib.testing.decorators import check_figures_equal
+import matplotlib.units as munits
 
 
 class TestUnitData:
@@ -321,3 +322,13 @@ def test_set_lim():
     ax.plot(["a", "b", "c", "d"], [1, 2, 3, 4])
     with warnings.catch_warnings():
         ax.set_xlim("b", "c")
+
+
+def test_mappable_error():
+    data = data = [["one", "two"], ["three", "four"]]
+
+    fig, ax = plt.subplots()
+    with pytest.raises(munits.ConversionError,
+                       match=('Categorical data does not support unit conversion '
+                              'that is not attached to an x or y Axis.')):
+        ax.imshow(data)
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -314,9 +314,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")
