ENH: expose max_labels in ColorMap (#90)

* ENH: expose max_lables

* lint

* docstring

* revert addition of the last value

* tests

* better error handling

Author: martinfleis

branca/colormap.py

@@ -70,24 +70,27 @@ class ColorMap(MacroElement):
         The right bound of the color scale.
     caption: str
         A caption to draw with the colormap.
+    max_labels : int, default 10
+        Maximum number of legend tick labels
     """
     _template = ENV.get_template('color_scale.js')
 
-    def __init__(self, vmin=0., vmax=1., caption=''):
+    def __init__(self, vmin=0., vmax=1., caption='', max_labels=10):
         super(ColorMap, self).__init__()
         self._name = 'ColorMap'
 
         self.vmin = vmin
         self.vmax = vmax
         self.caption = caption
         self.index = [vmin, vmax]
+        self.max_labels = max_labels
 
     def render(self, **kwargs):
         """Renders the HTML representation of the element."""
         self.color_domain = [self.vmin + (self.vmax-self.vmin) * k/499. for
                              k in range(500)]
         self.color_range = [self.__call__(x) for x in self.color_domain]
-        self.tick_labels = legend_scaler(self.index)
+        self.tick_labels = legend_scaler(self.index, self.max_labels)
 
         super(ColorMap, self).render(**kwargs)
 
@@ -180,11 +183,13 @@ class LinearColormap(ColorMap):
         Values lower than `vmin` will be bound directly to `colors[0]`.
     vmax : float, default 1.
         The maximal value for the colormap.
-        Values higher than `vmax` will be bound directly to `colors[-1]`."""
+        Values higher than `vmax` will be bound directly to `colors[-1]`.
+    max_labels : int, default 10
+        Maximum number of legend tick labels"""
 
-    def __init__(self, colors, index=None, vmin=0., vmax=1., caption=''):
+    def __init__(self, colors, index=None, vmin=0., vmax=1., caption='', max_labels=10):
         super(LinearColormap, self).__init__(vmin=vmin, vmax=vmax,
-                                             caption=caption)
+                                             caption=caption, max_labels=max_labels)
 
         n = len(colors)
         if n < 2:
@@ -216,7 +221,7 @@ def rgba_floats_tuple(self, x):
                      in range(4))
 
     def to_step(self, n=None, index=None, data=None, method=None,
-                quantiles=None, round_method=None):
+                quantiles=None, round_method=None, max_labels=10):
         """Splits the LinearColormap into a StepColormap.
 
         Parameters
@@ -243,6 +248,8 @@ def to_step(self, n=None, index=None, data=None, method=None,
             * If 'log10', all values will be rounded to the nearest
             order-of-magnitude integer. For example, 2100 is rounded to
             2000, 2790 to 3000.
+        max_labels : int, default 10
+            Maximum number of legend tick labels
 
         Returns
         -------
@@ -324,9 +331,10 @@ def to_step(self, n=None, index=None, data=None, method=None,
 
         caption = self.caption
 
-        return StepColormap(colors, index=index, vmin=index[0], vmax=index[-1], caption=caption)
+        return StepColormap(colors, index=index, vmin=index[0], vmax=index[-1], caption=caption,
+                            max_labels=max_labels)
 
-    def scale(self, vmin=0., vmax=1.):
+    def scale(self, vmin=0., vmax=1., max_labels=10):
         """Transforms the colorscale so that the minimal and maximal values
         fit the given parameters.
         """
@@ -336,6 +344,7 @@ def scale(self, vmin=0., vmax=1.):
             vmin=vmin,
             vmax=vmax,
             caption=self.caption,
+            max_labels=max_labels
             )
 
 
@@ -364,11 +373,13 @@ class StepColormap(ColorMap):
     vmax : float, default 1.
         The maximal value for the colormap.
         Values higher than `vmax` will be bound directly to `colors[-1]`.
+    max_labels : int, default 10
+        Maximum number of legend tick labels
 
     """
-    def __init__(self, colors, index=None, vmin=0., vmax=1., caption=''):
+    def __init__(self, colors, index=None, vmin=0., vmax=1., caption='', max_labels=10):
         super(StepColormap, self).__init__(vmin=vmin, vmax=vmax,
-                                           caption=caption)
+                                           caption=caption, max_labels=max_labels)
 
         n = len(colors)
         if n < 1:
@@ -393,7 +404,7 @@ def rgba_floats_tuple(self, x):
         i = len([u for u in self.index if u < x])  # 0 < i < n.
         return tuple(self.colors[i-1])
 
-    def to_linear(self, index=None):
+    def to_linear(self, index=None, max_labels=10):
         """
         Transforms the StepColormap into a LinearColormap.
 
@@ -403,6 +414,8 @@ def to_linear(self, index=None):
                 The values corresponding to each color in the output colormap.
                 It has to be sorted.
                 If None, a regular grid between `vmin` and `vmax` is created.
+        max_labels : int, default 10
+            Maximum number of legend tick labels
 
         """
         if index is None:
@@ -412,9 +425,9 @@ def to_linear(self, index=None):
 
         colors = [self.rgba_floats_tuple(x) for x in index]
         return LinearColormap(colors, index=index,
-                              vmin=self.vmin, vmax=self.vmax)
+                              vmin=self.vmin, vmax=self.vmax, max_labels=max_labels)
 
-    def scale(self, vmin=0., vmax=1.):
+    def scale(self, vmin=0., vmax=1., max_labels=10):
         """Transforms the colorscale so that the minimal and maximal values
         fit the given parameters.
         """
@@ -424,6 +437,7 @@ def scale(self, vmin=0., vmax=1.):
             vmin=vmin,
             vmax=vmax,
             caption=self.caption,
+            max_labels=max_labels
             )
 
 

tests/test_colormap.py

@@ -4,6 +4,7 @@
 ----------------------
 """
 import branca.colormap as cm
+import pytest
 
 
 def test_simple_step():
@@ -55,3 +56,30 @@ def test_step_object():
     cm.step.PuBu_06.to_linear()
     cm.step.YlGn_06.scale(3, 12)
     cm.step._repr_html_()
+
+@pytest.mark.parametrize("max_labels,expected", [
+    (10, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]),
+    (5, [0.0, '', 2.0, '', 4.0, '', 6.0, '', 8.0, '']),
+    (3, [0.0, '', '', '', 4.0, '', '', '', 8.0, '', '', '']),
+])
+def test_max_labels_linear(max_labels, expected):
+    colorbar = cm.LinearColormap(['red'] * 10, vmin=0, vmax=9, max_labels=max_labels)
+    try:
+        colorbar.render()
+    except AssertionError: # rendering outside parent Figure raises error
+        pass
+    assert colorbar.tick_labels == expected
+
+
+@pytest.mark.parametrize("max_labels,expected", [
+    (10, [0.0, '', 2.0, '', 4.0, '', 6.0, '', 8.0, '', 10.0, '']),
+    (5, [0.0, '', '', 3.0, '', '', 6.0, '', '', 9.0, '', '']),
+    (3, [0.0, '', '', '', 4.0, '', '', '', 8.0, '', '', '']),
+])
+def test_max_labels_step(max_labels, expected):
+    colorbar = cm.StepColormap(['red', 'blue'] * 5, vmin=0, vmax=10, max_labels=max_labels)
+    try:
+        colorbar.render()
+    except AssertionError: # rendering outside parent Figure raises error
+        pass
+    assert colorbar.tick_labels == expected