Merge pull request #11 from SylvainCorlay/DataFrame

Adding a DataFrame trait type

Author: SylvainCorlay

setup.py

@@ -78,7 +78,8 @@
 
 install_requires = setuptools_args['install_requires'] = [
     'traitlets>=4.2.2',
-    'numpy'
+    'numpy',
+    'pandas'
 ]
 
 extras_require = setuptools_args['extras_require'] = {

traittypes/tests/test_traittypes.py

@@ -7,8 +7,9 @@
 from unittest import TestCase
 from traitlets import HasTraits, TraitError, observe
 from traitlets.tests.test_traitlets import TraitTestBase
-from traittypes import Array
+from traittypes import Array, DataFrame, Series
 import numpy as np
+import pandas as pd
 
 
 # Good / Bad value trait test cases
@@ -39,7 +40,7 @@ class TestArray(TestCase):
     def test_array_equal(self):
         notifications = []
         class Foo(HasTraits):
-            bar = Array(default_value=[1, 2])
+            bar = Array([1, 2])
             @observe('bar')
             def _(self, change):
                 notifications.append(change)
@@ -103,3 +104,72 @@ class Foo(HasTraits):
         foo.bar = new_value
         self.assertTrue(np.array_equal(foo.bar, new_value))
 
+
+class TestDataFrame(TestCase):
+
+    def test_df_equal(self):
+        notifications = []
+        class Foo(HasTraits):
+            bar = DataFrame([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 = DataFrame()
+            b = DataFrame(None, allow_none=True)
+            c = DataFrame([])
+        foo = Foo()
+        self.assertTrue(foo.a.equals(pd.DataFrame()))
+        self.assertTrue(foo.b is None)
+        self.assertTrue(foo.c.equals(pd.DataFrame([])))
+
+    def test_allow_none(self):
+        class Foo(HasTraits):
+            bar = DataFrame()
+            baz = DataFrame(allow_none=True)
+        foo = Foo()
+        with self.assertRaises(TraitError):
+            foo.bar = None
+        foo.baz = None
+
+
+class TestSeries(TestCase):
+
+    def test_sereis_equal(self):
+        notifications = []
+        class Foo(HasTraits):
+            bar = Series([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 = Series()
+            b = Series(None, allow_none=True)
+            c = Series([])
+        foo = Foo()
+        self.assertTrue(foo.a.equals(pd.Series()))
+        self.assertTrue(foo.b is None)
+        self.assertTrue(foo.c.equals(pd.Series([])))
+
+    def test_allow_none(self):
+        class Foo(HasTraits):
+            bar = Series()
+            baz = Series(allow_none=True)
+        foo = Foo()
+        with self.assertRaises(TraitError):
+            foo.bar = None
+        foo.baz = None

traittypes/traittypes.py

@@ -1,8 +1,50 @@
 from traitlets import TraitType, TraitError, Undefined
 import numpy as np
+import pandas as pd
 
 
-class Array(TraitType):
+class SciType(TraitType):
+
+    """A base traittype for numpy arrays, pandas dataframes and series."""
+
+    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
+
+
+class Array(SciType):
 
     """A numpy array trait type."""
 
@@ -36,38 +78,70 @@ def __init__(self, default_value=Undefined, allow_none=False, dtype=None, **kwar
         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
+class DataFrame(SciType):
 
-        Validators return the (potentially modified) value, which is either
-        assigned to the HasTraits attribute or input into the next validator.
+    """A pandas dataframe trait type."""
 
-        They are evaluated in the order in which they are provided to the `valid`
-        function.
+    info_text = 'a pandas dataframe'
 
-        Example
-        -------
+    def validate(self, obj, value):
+        if value is None and not self.allow_none:
+            self.error(obj, value)
+        try:
+            value = pd.DataFrame(value)
+            for validator in self.validators:
+                value = validator(self, value)
+            return value
+        except (ValueError, TypeError) as e:
+            raise TraitError(e)
 
-        .. 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
+    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)
 
-            class Foo(HasTraits):
-                bar = Array(np.identity(2)).valid(shape(2, 2))
-            foo = Foo()
+    def __init__(self, default_value=Undefined, allow_none=False, dtype=None, **kwargs):
+        self.dtype = dtype
+        if default_value is Undefined:
+            default_value = pd.DataFrame()
+        elif default_value is not None:
+            default_value = pd.DataFrame(default_value)
+        self.validators = []
+        super(DataFrame, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)
 
-            foo.bar = [1, 2]  # Should raise a TraitError
-        """
-        self.validators.extend(validators)
-        return self
+
+class Series(SciType):
+
+    """A pandas series trait type."""
+
+    info_text = 'a pandas series'
+
+    def validate(self, obj, value):
+        if value is None and not self.allow_none:
+            self.error(obj, value)
+        try:
+            value = pd.Series(value)
+            for validator in self.validators:
+                value = validator(self, value)
+            return value
+        except (ValueError, TypeError) as e:
+            raise TraitError(e)
+
+    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):
+        self.dtype = dtype
+        if default_value is Undefined:
+            default_value = pd.Series()
+        elif default_value is not None:
+            default_value = pd.Series(default_value)
+        self.validators = []
+        super(Series, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)