From df2b3cd0b9d0741bbd899518aeadd2b4b55afe2d Mon Sep 17 00:00:00 2001
From: Hassan Kibirige <has2k1@gmail.com>
Date: Fri, 9 Mar 2018 12:22:51 +0300
Subject: [PATCH] Add cmap_d_pal

---
 doc/changelog.rst             |  1 +
 mizani/palettes.py            | 57 +++++++++++++++++++++++++++++++++--
 mizani/tests/test_palettes.py | 20 ++++++++++++
 3 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index a6cb7a3..d4df7de 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -6,6 +6,7 @@ v0.4.5
 *(not-yet-released)*
 
 - Added :class:`~mizani.palettes.identity_pal`
+- Added :class:`~mizani.palettes.cmap_d_pal`
 
 v0.4.4
 ------
diff --git a/mizani/palettes.py b/mizani/palettes.py
index 4399a89..ba6c75c 100644
--- a/mizani/palettes.py
+++ b/mizani/palettes.py
@@ -31,6 +31,7 @@
 __all__ = ['hls_palette', 'husl_palette', 'rescale_pal',
            'area_pal', 'abs_area', 'grey_pal', 'hue_pal',
            'brewer_pal', 'gradient_n_pal', 'cmap_pal',
+           'cmap_d_pal',
            'desaturate_pal', 'manual_pal', 'xkcd_palette',
            'crayon_palette', 'cubehelix_pal']
 
@@ -284,7 +285,7 @@ def hue_pal(h=.01, l=.6, s=.65, color_space='hls'):
     Returns
     -------
     out : function
-        Continuous color palette that takes a single
+        A discrete color palette that takes a single
         :class:`int` parameter ``n`` and returns ``n``
         equally spaced colors. Though the palette
         is continuous, since it is varies the hue it
@@ -517,7 +518,7 @@ def _gradient_n_pal(vals):
 
 def cmap_pal(name=None, lut=None):
     """
-    Create a palette using an MPL colormap
+    Create a continuous palette using an MPL colormap
 
     Parameters
     ----------
@@ -550,6 +551,58 @@ def _cmap_pal(vals):
     return _cmap_pal
 
 
+def cmap_d_pal(name=None, lut=None):
+    """
+    Create a discrete palette using an MPL Listed colormap
+
+    Parameters
+    ----------
+    name : str
+        Name of colormap
+    lut : None | int
+        This is the number of entries desired in the lookup table.
+        Default is ``None``, leave it up Matplotlib.
+
+    Returns
+    -------
+    out : function
+        A discrete color palette that takes a single
+        :class:`int` parameter ``n`` and returns ``n``
+        colors. The maximum value of ``n`` varies
+        depending on the parameters.
+
+    Examples
+    --------
+    >>> palette = cmap_d_pal('viridis')
+    >>> palette(5)
+    ['#440154', '#3b528b', '#21918c', '#5cc863', '#fde725']
+    """
+    colormap = get_cmap(name, lut)
+
+    if not isinstance(colormap, mcolors.ListedColormap):
+        raise ValueError(
+            "For a discrete palette, cmap must be of type "
+            "matplotlib.colors.ListedColormap")
+
+    ncolors = len(colormap.colors)
+
+    def _cmap_d_pal(n):
+        if n > ncolors:
+            raise ValueError(
+                "cmap `{}` has {} colors you requested {} "
+                "colors.".format(name, ncolors, n))
+
+        if ncolors < 256:
+            return [mcolors.rgb2hex(c) for c in colormap.colors[:n]]
+        else:
+            # Assume these are continuous and get colors equally spaced
+            # intervals  e.g. viridis is defined with 256 colors
+            idx = np.linspace(0, ncolors-1, n).round().astype(int)
+            return [mcolors.rgb2hex(colormap.colors[i]) for i in idx]
+
+    return _cmap_d_pal
+
+
 def desaturate_pal(color, prop, reverse=False):
     """
     Create a palette that desaturate a color by some proportion
diff --git a/mizani/tests/test_palettes.py b/mizani/tests/test_palettes.py
index e5032ab..5dc1c72 100644
--- a/mizani/tests/test_palettes.py
+++ b/mizani/tests/test_palettes.py
@@ -8,6 +8,7 @@
 from mizani.palettes import (hls_palette, husl_palette, rescale_pal,
                              area_pal, abs_area, grey_pal, hue_pal,
                              brewer_pal, gradient_n_pal, cmap_pal,
+                             cmap_d_pal,
                              desaturate_pal, manual_pal, xkcd_palette,
                              crayon_palette, cubehelix_pal, identity_pal)
 
@@ -119,6 +120,25 @@ def test_cmap_pal():
     assert all(s[0] == '#' and len(s) == 7 for s in result)
 
 
+def test_cmap_d_pal():
+    palette = cmap_d_pal('viridis')
+    result = palette(6)
+    assert all(s[0] == '#' and len(s) == 7 for s in result)
+    assert len(result) == 6
+
+    # More colors than palette
+    palette = cmap_d_pal('Accent')
+    result = palette(5)
+    assert all(s[0] == '#' and len(s) == 7 for s in result)
+
+    with pytest.raises(ValueError):
+        result = palette(10)
+
+    # Bad palette
+    with pytest.raises(ValueError):
+        palette = cmap_d_pal('Greens')
+
+
 def test_desaturate_pal():
     x = [0, .25, .5, .75, 1]
     # When desaturating pure red, green and blue