Merge pull request #23 from vidartf/tweaks

Further tweak definitions

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, TraitError, observe
+from traitlets import HasTraits, TraitError, observe, Undefined
 from traitlets.tests.test_traitlets import TraitTestBase
 from traittypes import Array, DataFrame, Series
 import numpy as np
@@ -56,11 +56,13 @@ class Foo(HasTraits):
             b = Array(dtype='int')
             c = Array(None, allow_none=True)
             d = Array([])
+            e = Array(Undefined)
         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, []))
+        self.assertTrue(foo.e is Undefined)
 
     def test_allow_none(self):
         class Foo(HasTraits):
@@ -116,19 +118,21 @@ def _(self, change):
                 notifications.append(change)
         foo = Foo()
         foo.bar = [1, 2]
-        self.assertFalse(len(notifications))
+        self.assertEqual(notifications, [])
         foo.bar = [1, 1]
-        self.assertTrue(len(notifications))
+        self.assertEqual(len(notifications), 1)
 
     def test_initial_values(self):
         class Foo(HasTraits):
             a = DataFrame()
             b = DataFrame(None, allow_none=True)
             c = DataFrame([])
+            d = DataFrame(Undefined)
         foo = Foo()
         self.assertTrue(foo.a.equals(pd.DataFrame()))
         self.assertTrue(foo.b is None)
         self.assertTrue(foo.c.equals(pd.DataFrame([])))
+        self.assertTrue(foo.d is Undefined)
 
     def test_allow_none(self):
         class Foo(HasTraits):
@@ -142,7 +146,7 @@ class Foo(HasTraits):
 
 class TestSeries(TestCase):
 
-    def test_sereis_equal(self):
+    def test_series_equal(self):
         notifications = []
         class Foo(HasTraits):
             bar = Series([1, 2])
@@ -151,19 +155,21 @@ def _(self, change):
                 notifications.append(change)
         foo = Foo()
         foo.bar = [1, 2]
-        self.assertFalse(len(notifications))
+        self.assertEqual(notifications, [])
         foo.bar = [1, 1]
-        self.assertTrue(len(notifications))
+        self.assertEqual(len(notifications), 1)
 
     def test_initial_values(self):
         class Foo(HasTraits):
             a = Series()
             b = Series(None, allow_none=True)
             c = Series([])
+            d = Series(Undefined)
         foo = Foo()
         self.assertTrue(foo.a.equals(pd.Series()))
         self.assertTrue(foo.b is None)
         self.assertTrue(foo.c.equals(pd.Series([])))
+        self.assertTrue(foo.d is Undefined)
 
     def test_allow_none(self):
         class Foo(HasTraits):

traittypes/tests/test_validators.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+import pytest
+
+from traitlets import HasTraits, TraitError
+
+from ..traittypes import SciType
+
+
+def test_coercion_validator():
+    # Test with a squeeze coercion
+    def truncate(trait, value):
+        return value[:10]
+
+    class Foo(HasTraits):
+        bar = SciType().valid(truncate)
+
+    foo = Foo(bar=list(range(20)))
+    assert foo.bar == list(range(10))
+    foo.bar = list(range(10, 40))
+    assert foo.bar == list(range(10, 20))
+
+
+def test_validaton_error():
+    # Test with a squeeze coercion
+    def maxlen(trait, value):
+        if len(value) > 10:
+            raise ValueError('Too long sequence!')
+        return value
+
+    class Foo(HasTraits):
+        bar = SciType().valid(maxlen)
+
+    # Check that it works as expected:
+    foo = Foo(bar=list(range(5)))
+    assert foo.bar == list(range(5))
+    # Check that it fails as expected:
+    with pytest.raises(TraitError):  # Should convert ValueError to TraitError
+        foo.bar = list(range(10, 40))
+    assert foo.bar == list(range(5))
+    # Check that it can again be set correctly
+    foo = Foo(bar=list(range(5, 10)))
+    assert foo.bar == list(range(5, 10))

traittypes/traittypes.py

@@ -1,6 +1,7 @@
+import inspect
 import warnings
 
-from traitlets import TraitType, TraitError, Undefined
+from traitlets import TraitType, TraitError, Undefined, Sentinel
 
 class _DelayedImportError(object):
     def __init__(self, package_name):
@@ -20,6 +21,13 @@ def __getattribute__(self, name):
     pd = _DelayedImportError('pandas')
 
 
+Empty = Sentinel('Empty', 'traittypes',
+"""
+Used in traittypes to specify that the default value should
+be an empty dataset
+""")
+
+
 class SciType(TraitType):
 
     """A base trait type for numpy arrays, pandas dataframes and series."""
@@ -107,96 +115,94 @@ 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=Empty, allow_none=False, dtype=None, **kwargs):
         self.dtype = dtype
-        if default_value is Undefined:
+        if default_value is Empty:
             default_value = np.array(0, dtype=self.dtype)
-        elif default_value is not None:
+        elif default_value is not None and default_value is not Undefined:
             default_value = np.asarray(default_value, dtype=self.dtype)
         super(Array, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)
 
     def make_dynamic_default(self):
-        if self.default_value is None:
+        if self.default_value is None or self.default_value is Undefined:
             return self.default_value
         else:
             return np.copy(self.default_value)
 
 
-class DataFrame(SciType):
+class PandasType(SciType):
 
     """A pandas dataframe trait type."""
 
     info_text = 'a pandas dataframe'
 
+    klass = None
+
     def validate(self, obj, value):
         if value is None and not self.allow_none:
             self.error(obj, value)
         if value is None or value is Undefined:
-            return super(DataFrame, self).validate(obj, value)
+            return super(PandasType, self).validate(obj, value)
         try:
-            value = pd.DataFrame(value)
+            value = self.klass(value)
         except (ValueError, TypeError) as e:
             raise TraitError(e)
-        return super(DataFrame, self).validate(obj, value)
+        return super(PandasType, self).validate(obj, value)
 
     def set(self, obj, value):
         new_value = self._validate(obj, value)
         old_value = obj._trait_values.get(self.name, self.default_value)
         obj._trait_values[self.name] = new_value
-        if (old_value is None and new_value is not None) or not old_value.equals(new_value):
+        if ((old_value is None and new_value is not None) or
+                (old_value is Undefined and new_value is not Undefined) or
+                not old_value.equals(new_value)):
             obj._notify_trait(self.name, old_value, new_value)
 
-    def __init__(self, default_value=Undefined, allow_none=False, dtype=None, **kwargs):
-        import pandas as pd
+    def __init__(self, default_value=Empty, allow_none=False, dtype=None, klass=None, **kwargs):
+        if klass is None:
+            klass = self.klass
+        if (klass is not None) and inspect.isclass(klass):
+            self.klass = klass
+        else:
+            raise TraitError('The klass attribute must be a class'
+                                ' not: %r' % klass)
         self.dtype = dtype
-        if default_value is Undefined:
-            default_value = pd.DataFrame()
-        elif default_value is not None:
-            default_value = pd.DataFrame(default_value)
-        super(DataFrame, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)
+        if default_value is Empty:
+            default_value = klass()
+        elif default_value is not None and default_value is not Undefined:
+            default_value = klass(default_value)
+        super(PandasType, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)
 
     def make_dynamic_default(self):
-        if self.default_value is None:
+        if self.default_value is None or self.default_value is Undefined:
             return self.default_value
         else:
             return self.default_value.copy()
 
 
-class Series(SciType):
+class DataFrame(PandasType):
 
-    """A pandas series trait type."""
+    """A pandas dataframe trait type."""
 
-    info_text = 'a pandas series'
+    info_text = 'a pandas dataframe'
 
-    def validate(self, obj, value):
-        if value is None and not self.allow_none:
-            self.error(obj, value)
-        if value is None or value is Undefined:
-            return super(Series, self).validate(obj, value)
-        try:
-            value = pd.Series(value)
-        except (ValueError, TypeError) as e:
-            raise TraitError(e)
-        return super(Series, self).validate(obj, value)
+    def __init__(self, default_value=Empty, allow_none=False, dtype=None, **kwargs):
+        if 'klass' not in kwargs and self.klass is None:
+            import pandas as pd
+            kwargs['klass'] = pd.DataFrame
+        super(DataFrame, self).__init__(
+            default_value=default_value, allow_none=allow_none, dtype=dtype, **kwargs)
 
-    def set(self, obj, value):
-        new_value = self._validate(obj, value)
-        old_value = obj._trait_values.get(self.name, self.default_value)
-        obj._trait_values[self.name] = new_value
-        if (old_value is None and new_value is not None) or not old_value.equals(new_value):
-            obj._notify_trait(self.name, old_value, new_value)
 
-    def __init__(self, default_value=Undefined, allow_none=False, dtype=None, **kwargs):
-        import pandas as pd
-        self.dtype = dtype
-        if default_value is Undefined:
-            default_value = pd.Series()
-        elif default_value is not None:
-            default_value = pd.Series(default_value)
-        super(Series, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)
+class Series(PandasType):
 
-    def make_dynamic_default(self):
-        if self.default_value is None:
-            return self.default_value
-        else:
-            return self.default_value.copy()
+    """A pandas series trait type."""
+
+    info_text = 'a pandas series'
+
+    def __init__(self, default_value=Empty, allow_none=False, dtype=None, **kwargs):
+        if 'klass' not in kwargs and self.klass is None:
+            import pandas as pd
+            kwargs['klass'] = pd.Series
+        super(Series, self).__init__(
+            default_value=default_value, allow_none=allow_none, dtype=dtype, **kwargs)