exploring xarray accessors

Author: Joseph Hamman

xbatcher/__init__.py

@@ -1 +1,2 @@
-from . generators import BatchGenerator
+from .generators import BatchGenerator
+from .accessors import BatchAccessor

xbatcher/accessors.py

@@ -0,0 +1,13 @@
+import xarray as xr
+
+from .generators import BatchGenerator
+
+
+@xr.register_dataarray_accessor("batch")
+@xr.register_dataset_accessor("batch")
+class BatchAccessor:
+    def __init__(self, xarray_obj):
+        self._obj = xarray_obj
+
+    def generator(self, *args, **kwargs):
+        return BatchGenerator(self._obj, *args, **kwargs)

xbatcher/tests/test_accessors.py

@@ -0,0 +1,34 @@
+import xarray as xr
+import numpy as np
+import pytest
+
+from xbatcher import BatchGenerator
+import xbatcher
+
+
+@pytest.fixture(scope='module')
+def sample_ds_3d():
+    shape = (10, 50, 100)
+    ds = xr.Dataset({'foo': (['time', 'y', 'x'], np.random.rand(*shape)),
+                     'bar': (['time', 'y', 'x'], np.random.randint(0, 10, shape))},
+                    {'x': (['x'], np.arange(shape[-1])),
+                     'y': (['y'], np.arange(shape[-2]))})
+    return ds
+
+
+def test_batch_accessor_ds(sample_ds_3d):
+    bg_class = BatchGenerator(sample_ds_3d, input_dims={'x': 5})
+    bg_acc = sample_ds_3d.batch.generator(input_dims={'x': 5})
+    assert isinstance(bg_acc, BatchGenerator)
+    for batch_class, batch_acc in zip(bg_class, bg_acc):
+        assert isinstance(batch_acc, xr.Dataset)
+        assert batch_class.equals(batch_acc)
+
+
+def test_batch_accessor_da(sample_ds_3d):
+    sample_da = sample_ds_3d['foo']
+    bg_class = BatchGenerator(sample_da, input_dims={'x': 5})
+    bg_acc = sample_da.batch.generator(input_dims={'x': 5})
+    assert isinstance(bg_acc, BatchGenerator)
+    for batch_class, batch_acc in zip(bg_class, bg_acc):
+        assert batch_class.equals(batch_acc)