Updated hex color validator and test for invalid values #2

AndreLobato · rmorshea · commit d49884fa9df2 · 2015-08-14T15:28:39.000-07:00
diff --git a/lib/matplotlib/tests/test_traitlets.py b/lib/matplotlib/tests/test_traitlets.py
@@ -3,22 +3,24 @@
 from nose.tools import *
 from unittest import TestCase
 from matplotlib.mpl_traitlets import Color, HasTraits
+from IPython.utils.traitlets import TraitError
 
 class ColorTestCase(TestCase):
     """Tests for the Color traits"""
 
     def setUp(self):
         self.transparent_values = [None, False, '', 'none']
-        self.black_values = ['#000000', (0,0,0,0), 0, 0.0, (.0,.0,.0), (.0,.0,.0,.0)]
+        self.black_values = ['#000000', '#000',(0,0,0,0), 0, 0.0, (.0,.0,.0), (.0,.0,.0,.0)]
         self.colored_values = ['#BE3537', (190,53,55), (0.7451, 0.20784, 0.21569)]
-        self.invalid_values = ['áfaef', '#FFF', '#0SX#$S', (0,0,0), (0.45,0.3), (()), {}, True]
+        self.unvalid_values = ['áfaef', '#0SX#$S', (0.45,0.3), 3.4, 344, (()), {}, True]
 
     def _evaluate_unvalids(self, a):
-        for values in self.invalid_values:
+        for values in self.unvalid_values:
             try:
                 a.color = values
-            except:
-                assert_raises(TypeError)
+                assert_true(False)
+            except TraitError:
+                assert_raises(TraitError)
 
     def test_noargs(self):
         class A(HasTraits):
diff --git a/lib/matplotlib/traitlets.py b/lib/matplotlib/traitlets.py
@@ -13,6 +13,7 @@
 #                                      Union, TraitError, HasTraits, 
 #                                      NoDefaultSpecified, TraitType)
 import numpy as np
+import re
 
 # override for backward compatability
 class Configurable(Configurable): pass
@@ -84,6 +85,7 @@ class Color(TraitType):
     info_text = 'float, int, tuple of float or int, or a hex string color'
     default_value = (0.0,0.0,0.0,0.0)
     named_colors = {}
+    _is_hex16 = re.compile(r'#[a-fA-F0-9]{3}(?:[a-fA-F0-9]{3})?$')
 
     def _int_to_float(self, value):
         as_float = (np.array(value)/255).tolist()
@@ -99,18 +101,13 @@ def _int_to_hex(self, value):
         return as_hex
 
     def _hex_to_float(self, value):
-        # Expects #FFFFFF format
-        split_hex = (value[1:3],value[3:5],value[5:7])
-        as_float = (np.array([int(v,16) for v in split_hex])/255.0).tolist()
+        if len(value) == 7:
+            split_hex = (value[1:3],value[3:5],value[5:7])
+            as_float = (np.array([int(v,16) for v in split_hex])/255.0).tolist()
+        elif len(value) == 4:
+            as_float = (np.array([int(v+v,16) for v in value[1:]])/255.0).tolist()
         return as_float
 
-    def _is_hex16(self, value):
-        try:
-            int(value, 16)
-            return True
-        except:
-            return False
-
     def _float_to_shade(self, value):
         grade = value*255.0
         return (grade,grade,grade)
@@ -121,21 +118,28 @@ def _int_to_shade(self, value):
 
     def validate(self, obj, value):
         in_range = False
-        if value is None or value is False or value in ['none','']:
-            # Return transparent if no other default alpha was set
-            return (0.0, 0.0, 0.0, 1.0)
+        if value is True:
+            self.error(obj, value)
 
-        if isinstance(value, float) and 0 <= value <= 1:
-            value = self._float_to_shade(value)
-        else:
-            in_range = False
+        elif value is None or value is False or value in ['none','']:
+            value = (0.0, 0.0, 0.0, 1.0)
+            in_range = True
 
-        if isinstance(value, int) and 0 <= value <= 255:
-            value = self._int_to_shade(value)
-        else:
-            in_range = False
+        elif isinstance(value, float):
+            if 0 <= value <= 1:
+                value = self._float_to_shade(value)
+                in_range = True
+            else:
+                in_range = False
 
-        if isinstance(value, (tuple, list)) and len(value) in (3,4):
+        elif isinstance(value, int):
+            if 0 <= value <= 255:
+                value = self._int_to_shade(value)
+                in_range = True
+            else:
+                in_range = False
+
+        elif isinstance(value, (tuple, list)) and len(value) in (3,4):
             is_all_float = np.prod([isinstance(v, (float)) for v in value])
             in_range = np.prod([(0 <= v <= 1) for v in value])
             if is_all_float and in_range:
@@ -146,10 +150,8 @@ def validate(self, obj, value):
                 if is_all_int and in_range:
                     value = self._int_to_float(value)
 
-        if isinstance(value, str) and len(value) == 7 and value[0] == '#':
-            is_all_hex16 = np.prod([self._is_hex16(v) for v in\
-                                    (value[1:3],value[3:5],value[5:7])])
-            if is_all_hex16:
+        elif isinstance(value, str) and len(value) in [4,7] and value[0] == '#':
+            if self._is_hex16.match(value):
                 value = self._hex_to_float(value)
                 in_range = np.prod([(0 <= v <= 1) for v in value])
                 if in_range:
@@ -158,7 +160,7 @@ def validate(self, obj, value):
         elif isinstance(value, str) and value in self.named_colors:
             value = self.validate(obj, self.named_colors[value])
             in_range = True
-
+            
         if in_range:
             if self._metadata['as_hex']:
                 return self._float_to_hex(value)
