Add custom validation ot Array trait type

Author: SylvainCorlay

README.md

@@ -6,19 +6,63 @@ Trait types for NumPy, SciPy and friends
 
 ## Goals
 
-Provide a reference implementation of trait types for common data structures used in the scipy stack such as
+Provide a reference implementation of trait types for common data structures
+used in the scipy stack such as
  - numpy arrays
  - pandas / xray data structures
 
-which are out of the scope of the main [traitlets](https://github.com/ipython/traitlets) project but are a common requirement to build applications with traitlets in combination with the scipy stack.
+which are out of the scope of the main [traitlets](https://github.com/ipython/traitlets)
+project but are a common requirement to build applications with traitlets in
+combination with the scipy stack.
 
-Another goal is to create adequate serialization and deserialization routines for these trait types to be used with the [ipywidgets](https://github.com/ipython/ipywidgets) project (`to_json` and `from_json`). These could also return a list of binary buffers as allowed by the current message protocol. 
+Another goal is to create adequate serialization and deserialization routines
+for these trait types to be used with the [ipywidgets](https://github.com/ipython/ipywidgets)
+project (`to_json` and `from_json`). These could also return a list of binary
+buffers as allowed by the current messaging protocol.
 
 ## Installation
 
-For a local installation, make sure you have
-[pip installed](https://pip.readthedocs.org/en/stable/installing/) and run:
+
+Using `pip`:
+
+Make sure you have [pip installed](https://pip.readthedocs.org/en/stable/installing/) and run:
 
 ```
 pip install traittypes
 ```
+
+Using `conda`:
+
+```
+conda install -c conda-forge traittypes
+```
+
+## Usage
+
+The `Array` trait type provide an implementation of a trait type for the numpy
+array.
+ - `Array` overrides some methods from `TraiType` that are generally not
+overloaded in order to work around some limitations with numpy array
+comparison.
+ - `Array` provides an API for adding custom validators to constained proposed
+values for the attribute.
+
+```python
+from traitlets import HasTraits, TraitError
+from traittypes import Array
+
+def shape(*dimensions):
+    def validator(trait, value):
+        if value.shape != dimensions:
+            raise TraitError('Expected an of shape %s and got and array with shape %s' % (dimensions, value.shape))
+        else:
+            return value
+    return validator
+
+class Foo(HasTraits):
+    bar = Array(np.identity(2)).valid(shape(2, 2))
+foo = Foo()
+
+foo.bar = [1, 2]  # Should raise a TraitError
+```
+

traittypes/tests/test_traittypes.py

@@ -27,7 +27,6 @@ class TestIntArray(TraitTestBase):
     _good_values = [1, [1, 2, 3], [[1, 2, 3], [4, 5, 6]], np.array([1])]
     _bad_values = [[1, [0, 0]]]
 
-
     def assertEqual(self, v1, v2):
         return np.testing.assert_array_equal(v1, v2)
 
@@ -70,3 +69,37 @@ class Foo(HasTraits):
         with self.assertRaises(TraitError):
             foo.bar = None
         foo.baz = None
+
+    def test_custom_validators(self):
+        # Test with a squeeze coercion
+        def squeeze(trait, value):
+            if 1 in value.shape:
+                value = np.squeeze(value)
+            return value
+
+        class Foo(HasTraits):
+            bar = Array().valid(squeeze)
+
+        foo = Foo(bar=[[1], [2]])
+        self.assertTrue(np.array_equal(foo.bar, [1, 2]))
+        foo.bar = [[1], [2], [3]]
+        self.assertTrue(np.array_equal(foo.bar, [1, 2, 3]))
+
+        # Test with a shape constraint
+        def shape(*dimensions):
+            def validator(trait, value):
+                if value.shape != dimensions:
+                    raise TraitError('Expected an of shape %s and got and array with shape %s' % (dimensions, value.shape))
+                else:
+                    return value
+            return validator
+
+        class Foo(HasTraits):
+            bar = Array(np.identity(2)).valid(shape(2, 2))
+        foo = Foo()
+        with self.assertRaises(TraitError):
+            foo.bar = [1]
+        new_value = [[0, 1], [1, 0]]
+        foo.bar = new_value
+        self.assertTrue(np.array_equal(foo.bar, new_value))
+

traittypes/traittypes.py

@@ -13,7 +13,10 @@ def validate(self, obj, value):
         if value is None and not self.allow_none:
             self.error(obj, value)
         try:
-            return np.asarray(value, dtype=self.dtype)
+            value = np.asarray(value, dtype=self.dtype)
+            for validator in self.validators:
+                value = validator(self, value)
+            return value
         except (ValueError, TypeError) as e:
             raise TraitError(e)
 
@@ -24,11 +27,47 @@ def set(self, obj, value):
         if not np.array_equal(old_value, new_value):
             obj._notify_trait(self.name, old_value, new_value)
 
-    def __init__(self, default_value=Undefined, allow_none=False,
-                 dtype=None, **kwargs):
+    def __init__(self, default_value=Undefined, allow_none=False, dtype=None, **kwargs):
         self.dtype = dtype
         if default_value is Undefined:
             default_value = np.array(0, dtype=self.dtype)
         elif default_value is not None:
             default_value = np.asarray(default_value, dtype=self.dtype)
+        self.validators = []
         super(Array, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)
+
+    def valid(self, *validators):
+        """
+        Register new trait validators
+
+        Validators are functions that take two arguments.
+         - The trait instance
+         - The proposed value
+
+        Validators return the (potentially modified) value, which is either
+        assigned to the HasTraits attribute or input into the next validator.
+
+        They are evaluated in the order in which they are provided to the `valid`
+        function.
+
+        Example
+        -------
+
+        .. code-block:: python
+            # Test with a shape constraint
+            def shape(*dimensions):
+                def validator(trait, value):
+                    if value.shape != dimensions:
+                        raise TraitError('Expected an of shape %s and got and array with shape %s' % (dimensions, value.shape))
+                    else:
+                        return value
+                return validator
+
+            class Foo(HasTraits):
+                bar = Array(np.identity(2)).valid(shape(2, 2))
+            foo = Foo()
+
+            foo.bar = [1, 2]  # Should raise a TraitError
+        """
+        self.validators.extend(validators)
+        return self