Merge pull request #23785 from tacaswell/fix_cmap_missing_typestability

greglucas · web-flow · commit 25b39d41d352 · 2022-09-01T20:46:01.000-06:00
FIX: ensure type stability for missing cmaps in `set_cmap`
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -701,6 +701,8 @@ 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"]
+    # use check_in_list to ensure type stability of the exception raised by
+    # the internal usage of this (ValueError vs KeyError)
+    _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)
