Add xarray DataArray

Author: davidbrochart

README.md

@@ -40,7 +40,7 @@ conda install -c conda-forge traittypes
 
 ## Usage
 
-`traittypes` extends the `traitlets` library with an implementation of trait types for numpy arrays, pandas dataframes, pandas series, and xarray datasets.
+`traittypes` extends the `traitlets` library with an implementation of trait types for numpy arrays, pandas dataframes, pandas series, xarray datasets and xarray dataarrays.
  - `traittypes` works around some limitations with numpy array comparison to only trigger change events when necessary.
  - `traittypes` also extends the traitlets API for adding custom validators to constained proposed values for the attribute.
 

docs/source/api_documentation.rst

@@ -24,3 +24,7 @@ The ``Series`` trait type holds a pandas Series.
 The ``Dataset`` trait type holds an xarray Dataset.
 
 .. autoclass:: traittypes.traittypes.Dataset
+
+The ``DataArray`` trait type holds an xarray DataArray.
+
+.. autoclass:: traittypes.traittypes.DataArray

traittypes/tests/test_traittypes.py

@@ -7,7 +7,7 @@
 from unittest import TestCase
 from traitlets import HasTraits, TraitError, observe, Undefined
 from traitlets.tests.test_traitlets import TraitTestBase
-from traittypes import Array, DataFrame, Series, Dataset
+from traittypes import Array, DataFrame, Series, Dataset, DataArray
 import numpy as np
 import pandas as pd
 import xarray as xr
@@ -201,12 +201,10 @@ def test_initial_values(self):
         class Foo(HasTraits):
             a = Dataset()
             b = Dataset(None, allow_none=True)
-            c = Dataset([])
             d = Dataset(Undefined)
         foo = Foo()
         self.assertTrue(foo.a.equals(xr.Dataset()))
         self.assertTrue(foo.b is None)
-        self.assertTrue(foo.c.equals(xr.Dataset([])))
         self.assertTrue(foo.d is Undefined)
 
     def test_allow_none(self):
@@ -217,3 +215,29 @@ class Foo(HasTraits):
         with self.assertRaises(TraitError):
             foo.bar = None
         foo.baz = None
+
+
+class TestDataArray(TestCase):
+
+    def test_ds_equal(self):
+        notifications = []
+        class Foo(HasTraits):
+            bar = DataArray([[0, 1], [2, 3]])
+            @observe('bar')
+            def _(self, change):
+                notifications.append(change)
+        foo = Foo()
+        foo.bar = [[0, 1], [2, 3]]
+        self.assertEqual(notifications, [])
+        foo.bar = [[0, 1], [2, 4]]
+        self.assertEqual(len(notifications), 1)
+
+    def test_initial_values(self):
+        class Foo(HasTraits):
+            b = DataArray(None, allow_none=True)
+            c = DataArray([])
+            d = DataArray(Undefined)
+        foo = Foo()
+        self.assertTrue(foo.b is None)
+        self.assertTrue(foo.c.equals(xr.DataArray([])))
+        self.assertTrue(foo.d is Undefined)

traittypes/traittypes.py

@@ -26,7 +26,7 @@ def __getattribute__(self, name):
 
 class SciType(TraitType):
 
-    """A base trait type for numpy arrays, pandas dataframes, pandas series and xarray datasets."""
+    """A base trait type for numpy arrays, pandas dataframes, pandas series, xarray datasets and xarray dataarrays."""
 
     def __init__(self, **kwargs):
         super(SciType, self).__init__(**kwargs)
@@ -128,9 +128,9 @@ def make_dynamic_default(self):
 
 class PandasType(SciType):
 
-    """A pandas dataframe trait type."""
+    """A pandas dataframe or series trait type."""
 
-    info_text = 'a pandas dataframe'
+    info_text = 'a pandas dataframe or series'
 
     klass = None
 
@@ -154,15 +154,14 @@ def set(self, obj, value):
                 not old_value.equals(new_value)):
             obj._notify_trait(self.name, old_value, new_value)
 
-    def __init__(self, default_value=Empty, allow_none=False, dtype=None, klass=None, **kwargs):
+    def __init__(self, default_value=Empty, allow_none=False, 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 Empty:
             default_value = klass()
         elif default_value is not None and default_value is not Undefined:
@@ -195,20 +194,22 @@ class Series(PandasType):
     """A pandas series trait type."""
 
     info_text = 'a pandas series'
+    dtype = None
 
     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)
+        self.dtype = dtype
 
 
 class XarrayType(SciType):
 
-    """An xarray dataset trait type."""
+    """An xarray dataset or dataarray trait type."""
 
-    info_text = 'an xarray dataset'
+    info_text = 'an xarray dataset or dataarray'
 
     klass = None
 
@@ -232,15 +233,14 @@ def set(self, obj, value):
                 not old_value.equals(new_value)):
             obj._notify_trait(self.name, old_value, new_value)
 
-    def __init__(self, default_value=Empty, allow_none=False, dtype=None, klass=None, **kwargs):
+    def __init__(self, default_value=Empty, allow_none=False, 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 Empty:
             default_value = klass()
         elif default_value is not None and default_value is not Undefined:
@@ -266,3 +266,19 @@ def __init__(self, default_value=Empty, allow_none=False, dtype=None, **kwargs):
             kwargs['klass'] = xr.Dataset
         super(Dataset, self).__init__(
             default_value=default_value, allow_none=allow_none, dtype=dtype, **kwargs)
+
+
+class DataArray(XarrayType):
+
+    """An xarray dataarray trait type."""
+
+    info_text = 'an xarray dataarray'
+    dtype = None
+
+    def __init__(self, default_value=Empty, allow_none=False, dtype=None, **kwargs):
+        if 'klass' not in kwargs and self.klass is None:
+            import xarray as xr
+            kwargs['klass'] = xr.DataArray
+        super(DataArray, self).__init__(
+            default_value=default_value, allow_none=allow_none, dtype=dtype, **kwargs)
+        self.dtype = dtype