Custom colormap trait validation (#99)

* allow setting colormap on init

* flake8 fixes

* custom colormap trait validation

Author: akaszynski

docs/source/api_reference/isocolor.rst

@@ -3,9 +3,6 @@ IsoColor
 
 The ``IsoColor`` widget colorize your mesh given.
 
-.. note::
-    Currently, you can only use the default Viridis colormap, we'd like to support all paraview colormaps
-
 The ``input`` attribute should be the name of the ``Component`` you want to use for colorizing the mesh. You also need to pass the ``min`` and ``max`` (or ``range`` as a min/max tuple) of the component array.
 
 For example, if you have a 1-D ``Data`` named ``"height"``, you can simply pass its name as input:
@@ -73,8 +70,8 @@ Examples
     height_min = np.min(z)
     height_max = np.max(z)
 
-    # Colorize by height
-    colored_mesh = IsoColor(mesh, input='height', min=height_min, max=height_max)
+    # Colorize by height and set the colormap to "Turbo"
+    colored_mesh = IsoColor(mesh, input='height', min=height_min, max=height_max, colormap='Turbo')
 
     # Create a slider that will dynamically change the boundaries of the colormap
     colormap_slider_range = FloatRangeSlider(value=[height_min, height_max], min=height_min, max=height_max, step=(height_max - height_min) / 100.)

ipygany/ipygany.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from traitlets import (
-    Bool, Dict, Enum, Unicode, List, Instance, CFloat, Tuple, TraitError, Union, default, validate
+    Bool, Dict, Enum, Unicode, List, Instance, CFloat, Tuple, TraitError, Union, default, validate, Any,
 )
 from traittypes import Array
 from ipywidgets import (
@@ -647,7 +647,7 @@ class IsoColor(Effect):
     min = CFloat(0.).tag(sync=True)
     max = CFloat(0.).tag(sync=True)
     range = Tuple((0., 0.)).tag(sync=True)
-    colormap = Enum(list(colormaps.values()), allow_none=False, default_value=colormaps.Viridis).tag(sync=True)
+    colormap = Any(allow_none=False, default_value=colormaps.Viridis).tag(sync=True)
     type = Enum(['linear', 'log'], default_value='linear').tag(sync=True)
 
     def __init__(self, parent, **kwargs):
@@ -659,6 +659,31 @@ def input_dim(self):
         """Input dimension."""
         return 1
 
+    @validate('colormap')
+    def _valid_colormap(self, proposal):
+        """Validate colormap. """
+        # Logic is here to allow setting ``colormap`` using strings
+        # and ints.  See:
+        # https://github.com/QuantStack/ipygany/pull/99
+        value = proposal['value']
+
+        # set colormap
+        if isinstance(value, str):
+            if value not in colormaps:
+                allowed = ', '.join([f"'{key}'" for key in colormaps.keys()])
+                raise ValueError(f'``colormap`` "{value} is not supported by ``ipygany``\n'
+                                 'Pick from one of the following:\n' + allowed)
+            colormap = colormaps[value]
+        elif isinstance(value, int):
+            if value not in colormaps.values():
+                raise ValueError(f'Invalid colormap {value}.  Expected one of'
+                                 f'the following: {colormaps.values()}')
+            colormap = value
+        elif value is not None:
+            raise TypeError(f'Invalid colormap type {type(value)}.  Should be '
+                            'a valid colormap ``str`` or ``int``')
+        return colormap
+
 
 class ColorBar(_GanyDOMWidgetBase):
     """A ColorBar widget."""

tests/test_ipygany.py

@@ -1,5 +1,3 @@
-import pytest
-
 import numpy as np
 
 from ipydatawidgets import NDArrayWidget

tests/test_isocolor.py

@@ -1,8 +1,7 @@
 import pytest
 
-from traitlets import TraitError
-
 from ipygany import PolyMesh, IsoColor
+from ipygany.colormaps import colormaps
 
 from .utils import get_test_assets
 
@@ -12,7 +11,7 @@ def test_default_input():
 
     poly = PolyMesh(vertices=vertices, triangle_indices=triangles, data=[data_1d, data_3d])
 
-    colored_mesh = IsoColor(poly)
+    colored_mesh = IsoColor(poly, colormap='Turbo')
 
     assert colored_mesh.input == '1d'
 
@@ -23,6 +22,25 @@ def test_default_input():
     assert colored_mesh.input == (('3d', 'x'), )
 
 
+def test_colormaps():
+    vertices, triangles, data_1d, data_3d = get_test_assets()
+
+    poly = PolyMesh(vertices=vertices, triangle_indices=triangles, data=[data_1d, data_3d])
+
+    colored_mesh = IsoColor(poly, colormap='Turbo')
+    assert colored_mesh.colormap == colormaps.Turbo
+
+    colormap = colormaps.Cividis
+    colored_mesh.colormap = colormap
+    assert colored_mesh.colormap == colormap
+
+    with pytest.raises(ValueError):
+        colored_mesh.colormap = 'InvalidColor'
+
+    with pytest.raises(ValueError):
+        colored_mesh.colormap = -1
+
+
 def test_input():
     vertices, triangles, data_1d, data_3d = get_test_assets()