refactor: update colormap API for new versions of matplotlib

geemap/colormaps.py

@@ -240,6 +240,24 @@ def plot_colormaps(width: float = 8.0, height: float = 0.4) -> None:
     plt.show()
 
 
+def get_colormap(cmap_name: str, n_class: int | None = None) -> mpl.colors.Colormap:
+    """Returns a matplotlib colormap object"""
+    if hasattr(mpl, "colormaps"):
+        try:
+            cmap = mpl.colormaps[cmap_name]
+        except KeyError as e:
+            raise ValueError(f"Colormap '{cmap_name}' not found.") from e
+        if n_class is not None:
+            cmap = cmap.resampled(n_class)
+        return cmap
+    elif hasattr(plt, "get_cmap"):
+        cmap = plt.get_cmap(cmap_name, n_class)
+        return cmap
+    else:
+        cmap = plt.cm.get_cmap(cmap_name, n_class)
+    return cmap
+
+
 def get_palettes() -> box.Box:
     """Returns a dictionary of colormaps and their associated palettes."""
     for cmap_name in list_colormaps():

geemap/common.py

@@ -3656,7 +3656,7 @@ def save_colorbar(
         alpha = 1
 
     if cmap is not None:
-        cmap = mpl.pyplot.get_cmap(cmap)
+        cmap = colormaps.get_colormap(cmap)
         norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
 
     if "palette" in vis_params:
@@ -3673,7 +3673,7 @@ def save_colorbar(
             norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
 
     elif cmap is not None:
-        cmap = mpl.pyplot.get_cmap(cmap)
+        cmap = colormaps.get_colormap(cmap)
         norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
 
     else:
@@ -11919,10 +11919,9 @@ def classify(
 
     if cmap is None:
         cmap = "Blues"
-    try:
-        cmap = plt.get_cmap(cmap, k)
-    except:
-        cmap = plt.cm.get_cmap(cmap, k)
+
+    cmap = colormaps.get_colormap(cmap, k)
+
     if colors is None:
         colors = [mpl.colors.rgb2hex(cmap(i))[1:] for i in range(cmap.N)]
         colors = ["#" + i for i in colors]
@@ -12676,10 +12675,10 @@ def get_palette_colors(
         n_class: The number of colors. Defaults to None.
         hashtag: Whether to return a list of hex colors. Defaults to False.
     """
-    try:
-        cmap = plt.get_cmap(cmap_name, n_class)
-    except:
-        cmap = plt.cm.get_cmap(cmap_name, n_class)
+    if cmap_name is None:
+        cmap_name = "viridis"
+    cmap = colormaps.get_colormap(cmap_name, n_class)
+
     colors = [mpl.colors.rgb2hex(cmap(i))[1:] for i in range(cmap.N)]
     if hashtag:
         colors = ["#" + i for i in colors]