Accessors created with xr.register_*_accessor are cached

Author: shoyer

doc/internals.rst

@@ -78,7 +78,7 @@ xarray:
 
 .. literalinclude:: examples/_code/accessor_example.py
 
-This achieves the same result as if the ``Dataset`` class had a property
+This achieves the same result as if the ``Dataset`` class had a cached property
 defined that returns an instance of your class:
 
 .. python::
@@ -90,9 +90,14 @@ defined that returns an instance of your class:
           return GeoAccessor(self)
 
 However, using the register accessor decorators is preferable to simply adding
-your own ad-hoc property (i.e., ``Dataset.geo = property(...)``), because it
-ensures that your property does not conflict with any other attributes or
-methods.
+your own ad-hoc property (i.e., ``Dataset.geo = property(...)``), for two
+reasons:
+
+1. It ensures that the name of your property does not conflict with any other
+   attributes or methods.
+2. Instances of accessor object will be cached on the xarray object that creates
+   them. This means you can save state on them (e.g., to cache computed
+   properties).
 
 Back in an interactive IPython session, we can use these properties:
 
@@ -110,7 +115,7 @@ Back in an interactive IPython session, we can use these properties:
 
 The intent here is that libraries that extend xarray could add such an accessor
 to implement subclass specific functionality rather than using actual subclasses
-rather patching in a large number of domain specific methods.
+or patching in a large number of domain specific methods.
 
 To help users keep things straight, please `let us know
 <https://github.com/pydata/xarray/issues>`_ if you plan to write a new accessor

xarray/core/dataarray.py

@@ -1376,19 +1376,19 @@ def imag(self):
 
     def dot(self, other):
         """Perform dot product of two DataArrays along their shared dims.
-        
+
         Equivalent to taking taking tensordot over all shared dims.
 
         Parameters
         ----------
         other : DataArray
             The other array with which the dot product is performed.
-            
+
         Returns
         -------
         result : DataArray
             Array resulting from the dot product over all shared dimensions.
-            
+
         See also
         --------
         np.tensordot(a, b, axes)
@@ -1397,10 +1397,10 @@ def dot(self, other):
         --------
 
         >>> da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4))
-        >>> da = DataArray(da_vals, dims=['x', 'y', 'z'])    
+        >>> da = DataArray(da_vals, dims=['x', 'y', 'z'])
         >>> dm_vals = np.arange(4)
         >>> dm = DataArray(dm_vals, dims=['z'])
-                
+
         >>> dm.dims
         ('z')
         >>> da.dims

xarray/core/extensions.py

@@ -6,15 +6,33 @@ class AccessorRegistrationError(Exception):
     """Exception for conflicts in accessor registration."""
 
 
+class _CachedAccessor(object):
+    """Custom property-like object (descriptor) for caching accessors."""
+    def __init__(self, name, accessor):
+        self._name = name
+        self._accessor = accessor
+
+    def __get__(self, obj, cls):
+        if obj is None:
+            # we're accessing the attribute of the class, i.e., Dataset.geo
+            return self._accessor
+        accessor_obj = self._accessor(obj)
+        # Replace the property with the accessor object. Inspired by:
+        # http://www.pydanny.com/cached-property.html
+        # We need to use object.__setattr__ because we overwrite __setattr__ on
+        # AttrAccessMixin.
+        object.__setattr__(obj, self._name, accessor_obj)
+        return accessor_obj
+
+
 def _register_accessor(name, cls):
     def decorator(accessor):
         if hasattr(cls, name):
             raise AccessorRegistrationError(
                 'cannot register accessor %r under name %r for type %r '
                 'because an attribute with that name already exists.'
                 % (accessor, name, cls))
-
-        setattr(cls, name, property(accessor))
+        setattr(cls, name, _CachedAccessor(name, accessor))
         return accessor
     return decorator
 

xarray/test/test_extensions.py

@@ -1,14 +1,28 @@
+try:
+    import cPickle as pickle
+except ImportError:
+    import pickle
+
 import xarray as xr
 
 from . import TestCase
 
 
+@xr.register_dataset_accessor('example_accessor')
+@xr.register_dataarray_accessor('example_accessor')
+class ExampleAccessor(object):
+    """For the pickling tests below."""
+    def __init__(self, xarray_obj):
+        self.obj = xarray_obj
+
+
 class TestAccessor(TestCase):
     def test_register(self):
 
         @xr.register_dataset_accessor('demo')
         @xr.register_dataarray_accessor('demo')
         class DemoAccessor(object):
+            """Demo accessor."""
             def __init__(self, xarray_obj):
                 self._obj = xarray_obj
 
@@ -22,10 +36,41 @@ def foo(self):
         da = xr.DataArray(0)
         assert da.demo.foo == 'bar'
 
+        # accessor is cached
+        assert ds.demo is ds.demo
+
+        # check descriptor
+        assert ds.demo.__doc__ == "Demo accessor."
+        assert xr.Dataset.demo.__doc__ == "Demo accessor."
+        assert isinstance(ds.demo, DemoAccessor)
+        assert xr.Dataset.demo is DemoAccessor
+
+        # ensure we can remove it
         del xr.Dataset.demo
-        assert not hasattr(ds, 'demo')
+        assert not hasattr(xr.Dataset, 'demo')
 
         with self.assertRaises(xr.core.extensions.AccessorRegistrationError):
             @xr.register_dataarray_accessor('demo')
             class Foo(object):
                 pass
+
+        # it didn't get registered again
+        assert not hasattr(xr.Dataset, 'demo')
+
+    def test_pickle_dataset(self):
+        ds = xr.Dataset()
+        ds_restored = pickle.loads(pickle.dumps(ds))
+        assert ds.identical(ds_restored)
+
+        # state save on the accessor is restored
+        assert ds.example_accessor is ds.example_accessor
+        ds.example_accessor.value = 'foo'
+        ds_restored = pickle.loads(pickle.dumps(ds))
+        assert ds.identical(ds_restored)
+        assert ds_restored.example_accessor.value == 'foo'
+
+    def test_pickle_dataarray(self):
+        array = xr.Dataset()
+        assert array.example_accessor is array.example_accessor
+        array_restored = pickle.loads(pickle.dumps(array))
+        assert array.identical(array_restored)