provide cmap used on last plot

legendgram/legendgram.py

@@ -1,4 +1,5 @@
 from .util import make_location as _make_location
+from .util import _get_cmap
 import numpy as np
 from matplotlib.colors import Colormap
 import matplotlib.pyplot as plt
@@ -47,7 +48,7 @@ def legendgram(y, breaks=None, pal=None, bins=50, clip=None,
     if ax is None:
         ax = plt.gca()
     if pal is None:
-        pal = 'viridis'
+        pal = _get_cmap(ax)
     if breaks is None:
         breaks = 10
     if isinstance(breaks, int):

legendgram/util.py

@@ -1,3 +1,8 @@
+import warnings
+
+from matplotlib import colormaps as cm
+from matplotlib.axes._axes import Axes
+from matplotlib.colors import ListedColormap
 
 loc_lut = {'best'         : 0,
             'upper right'  : 1,
@@ -71,3 +76,23 @@ def make_location(ax,loc, legend_size=(.27,.2)):
     elif loc.lower() == 'upper right':
         anchor_x, anchor_y = position.x0 + right_offset, position.y0 + top_offset
     return [anchor_x, anchor_y, legend_width, legend_height]
+
+
+def _get_cmap(_ax: Axes) -> ListedColormap:
+    """Detect the most recent matplotlib colormap used, if previously rendered."""
+    _images = _ax.images
+    has_images = len(_images)
+    if has_images:
+        cmap = _images[-1].cmap
+        if has_images > 1:
+            warnings.warn(
+                (
+                    "More than one image associated with the axes. "
+                    f"Defaulting to last colormap: '{cmap.name}'"
+                ),
+                UserWarning,
+                stacklevel=2
+            )
+    else:
+        cmap = cm.get_cmap("viridis")
+    return cmap