From 2396aaac9b7cb66b514644caa4410dd5bb194e6d Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Tue, 10 May 2016 09:39:15 -0700 Subject: [PATCH] Accessors created with xr.register_*_accessor are cached --- doc/internals.rst | 13 +++++++--- xarray/core/dataarray.py | 10 ++++---- xarray/core/extensions.py | 22 ++++++++++++++-- xarray/test/test_extensions.py | 47 +++++++++++++++++++++++++++++++++- 4 files changed, 80 insertions(+), 12 deletions(-) diff --git a/doc/internals.rst b/doc/internals.rst index bb232fe128f..2d93932adc6 100644 --- a/doc/internals.rst +++ b/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: diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 47a8c463091..fc8a73335cb 100644 --- a/xarray/core/dataarray.py +++ b/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 diff --git a/xarray/core/extensions.py b/xarray/core/extensions.py index 7f7b7e9fb60..5df1e73868d 100644 --- a/xarray/core/extensions.py +++ b/xarray/core/extensions.py @@ -6,6 +6,25 @@ 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): @@ -13,8 +32,7 @@ def decorator(accessor): '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 diff --git a/xarray/test/test_extensions.py b/xarray/test/test_extensions.py index 465626dd8f3..7a751ea87bd 100644 --- a/xarray/test/test_extensions.py +++ b/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)