A convenient function to show a matplotlib color map

plot_cmap.py

import matplotlib.pyplot as plt

def plot_cmap(cmap, ncolor):
    """ 
    A convenient function to plot colors of a matplotlib cmap
    
    Args:
        ncolor (int): number of color to show
        cmap: a cmap object or a matplotlib color name
    """
    
    if isinstance(cmap, str):
        try:
            cm = plt.get_cmap(cmap)
        except ValueError:
            print("WARNINGS :", cmap, " is not a known colormap")
            cm = plt.cm.gray
    else:
        cm = cmap
            
    with mpl.rc_context(mpl.rcParamsDefault):
        fig = plt.figure(figsize=(6, 1), frameon=False)
        ax = fig.add_subplot(111)
        ax.pcolor(np.linspace(1, ncolor, ncolor).reshape(1, ncolor), cmap=cm)
        ax.set_title(cm.name)
        xt = ax.set_xticks([])
        yt = ax.set_yticks([])
    return fig

show_colors.py

def show_colors(colors):
    """
    Draw a square for each color contained in the colors list
    given in argument.
    """
    with plt.rc_context(plt.rcParamsDefault):
        fig = plt.figure(figsize=(6, 1), frameon=False)
        ax = fig.add_subplot(111)
        for x, color in enumerate(colors):
            ax.add_patch(
                mpl.patches.Rectangle(
                    (x, 0), 1, 1, facecolor=color
                )
            )
        ax.set_xlim((0, len(colors)))
        ax.set_ylim((0, 1))
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_aspect("equal")
    
    return fig

I got following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "plot_cmap.py", line 21, in plot_cmap
    with mpl.rc_context(mpl.rcParamsDefault):
NameError: name 'mpl' is not defined

I fixed it by adding

import matplotlib as mpl
import numpy as np