FIX: ensure type stability for missing cmaps in set_cmap

tacaswell · tacaswell · commit 0b6b385e6d44 · 2022-08-31T15:43:53.000-04:00
Previously this would raise a `ValueError`, the refactor caused it to raise
`KeyError`.   This inadvertently located a bad test in xarray, but restoring
the old behavior.
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -701,6 +701,6 @@ def _ensure_cmap(cmap):
     """
     if isinstance(cmap, colors.Colormap):
         return cmap
-    return mpl.colormaps[
-        cmap if cmap is not None else mpl.rcParams['image.cmap']
-    ]
+    cmap_name = cmap if cmap is not None else mpl.rcParams["image.cmap"]
+    _api.check_in_list(sorted(_colormaps), cmap=cmap_name)
+    return mpl.colormaps[cmap_name]
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1543,3 +1543,11 @@ def test_color_sequences():
     plt.color_sequences.unregister('rgb')  # multiple unregisters are ok
     with pytest.raises(ValueError, match="Cannot unregister builtin"):
         plt.color_sequences.unregister('tab10')
+
+
+def test_cm_set_cmap_error():
+    sm = cm.ScalarMappable()
+    # Pick a name we are pretty sure will never be a colormap name
+    bad_cmap = 'AardvarksAreAwkward'
+    with pytest.raises(ValueError, match=bad_cmap):
+        sm.set_cmap(bad_cmap)
