Merge pull request #6 from SylvainCorlay/default_handling

Better handling of default values

Author: SylvainCorlay

traittypes/tests/test_traittypes.py

@@ -5,7 +5,7 @@
 # Distributed under the terms of the Modified BSD License.
 
 from unittest import TestCase
-from traitlets import HasTraits, observe
+from traitlets import HasTraits, TraitError, observe
 from traitlets.tests.test_traitlets import TraitTestBase
 from traittypes import Array
 import numpy as np
@@ -39,15 +39,34 @@ class TestArray(TestCase):
 
     def test_array_equal(self):
         notifications = []
-
         class Foo(HasTraits):
             bar = Array(default_value=[1, 2])
             @observe('bar')
             def _(self, change):
                 notifications.append(change)
-
         foo = Foo()
         foo.bar = [1, 2]
         self.assertFalse(len(notifications))
         foo.bar = [1, 1]
         self.assertTrue(len(notifications))
+
+    def test_initial_values(self):
+        class Foo(HasTraits):
+            a = Array()
+            b = Array(dtype='int')
+            c = Array(None, allow_none=True)
+            d = Array([])
+        foo = Foo()
+        self.assertTrue(np.array_equal(foo.a, np.array(0)))
+        self.assertTrue(np.array_equal(foo.b, np.array(0)))
+        self.assertTrue(foo.c is None)
+        self.assertTrue(np.array_equal(foo.d, []))
+
+    def test_allow_none(self):
+        class Foo(HasTraits):
+            bar = Array()
+            baz = Array(allow_none=True)
+        foo = Foo()
+        with self.assertRaises(TraitError):
+            foo.bar = None
+        foo.baz = None

traittypes/traittypes.py

@@ -1,4 +1,4 @@
-from traitlets import TraitType, TraitError
+from traitlets import TraitType, TraitError, Undefined
 import numpy as np
 
 
@@ -7,11 +7,13 @@ class Array(TraitType):
     """A numpy array trait type."""
 
     info_text = 'a numpy array'
+    dtype = None
 
     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.get_metadata('dtype'),
-                              order=self.get_metadata('order'))
+            return np.asarray(value, dtype=self.dtype)
         except (ValueError, TypeError) as e:
             raise TraitError(e)
 
@@ -21,3 +23,12 @@ def set(self, obj, value):
         obj._trait_values[self.name] = new_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):
+        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)
+        super(Array, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)