Adding a Series dataframe

Author: SylvainCorlay

traittypes/tests/test_traittypes.py

@@ -7,7 +7,7 @@
 from unittest import TestCase
 from traitlets import HasTraits, TraitError, observe
 from traitlets.tests.test_traitlets import TraitTestBase
-from traittypes import Array, DataFrame
+from traittypes import Array, DataFrame, Series
 import numpy as np
 import pandas as pd
 
@@ -138,3 +138,38 @@ class Foo(HasTraits):
         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

@@ -111,3 +111,37 @@ def __init__(self, default_value=Undefined, allow_none=False, dtype=None, **kwar
             default_value = pd.DataFrame(default_value)
         self.validators = []
         super(DataFrame, self).__init__(default_value=default_value, allow_none=allow_none, **kwargs)
+
+
+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)