Merge pull request #23800 from meeseeksmachine/auto-backport-of-pr-23785-on-v3.6.x

QuLogic · web-flow · commit 2736f0b52634 · 2022-09-02T00:34:20.000-04:00
Backport PR #23785 on branch v3.6.x (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)
