Merge pull request #7 from jhamman/accessor

exploring xarray accessors

Author: Joe Hamman

doc/api.rst

@@ -1,8 +1,24 @@
+.. _api:
 
 API reference
 -------------
 
 This page provides an auto-generated summary of Xbatcher's API.
 
+Core
+====
+
 .. autoclass:: xbatcher.BatchGenerator
    :members:
+
+Dataset.batch and DataArray.batch
+=================================
+
+.. currentmodule:: xarray
+
+.. autosummary::
+   :toctree: generated/
+   :template: autosummary/accessor_method.rst
+
+   Dataset.batch.generator
+   DataArray.batch.generator

doc/conf.py

@@ -15,6 +15,8 @@
 import os
 import sys
 
+import sphinx_autosummary_accessors
+
 import xbatcher
 
 # If extensions (or modules to document with autodoc) are in another directory,
@@ -47,6 +49,7 @@
     'nbsphinx',
     'IPython.sphinxext.ipython_directive',
     'IPython.sphinxext.ipython_console_highlighting',
+    'sphinx_autosummary_accessors',
 ]
 
 # never execute notebooks: avoids lots of expensive imports on rtd
@@ -66,14 +69,16 @@ def setup(app):
 
 
 # link to github issues
-extlinks = {'issue': ('https://github.com/xbatcher/xbatcher/issues/%s', 'GH')}
+extlinks = {
+    'issue': ('https://github.com/pangeo-data/xbatcher/issues/%s', 'GH')
+}
 
 autosummary_generate = True
 numpydoc_class_members_toctree = False
 numpydoc_show_class_members = False
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ['_templates', sphinx_autosummary_accessors.templates_path]
 
 # The suffix of source filenames.
 source_suffix = '.rst'
@@ -86,7 +91,7 @@ def setup(app):
 
 # General information about the project.
 project = u'xbatcher'
-copyright = u'2016, xbatcher developers'
+copyright = u'2021, xbatcher developers'
 
 # The version info for the project you're documenting, acts as replacement for
 # |version| and |release|, also used in various other places throughout the

doc/environment.yml

@@ -1,14 +1,14 @@
 name: xbatcher
 dependencies:
-  - python=3.6
-  - xarray
+  - python=3.7
   - dask
+  - docrep
+  - ipython
+  - matplotlib
+  - nbsphinx
   - numpy
-  - pytest
   - numpydoc
+  - pytest
   - sphinx
-  - ipython
-  - matplotlib
-  - pip:
-      - docrep
-      - nbsphinx
+  - sphinx-autosummary-accessors
+  - xarray

doc/index.rst

@@ -39,6 +39,17 @@ and we want to create batches along the time dimension. We can do it like this
         # actually feed to machine learning library
     batch
 
+or via a built-in `Xarray accessor <http://xarray.pydata.org/en/stable/internals/extending-xarray.html#extending-xarray>`_:
+
+.. ipython:: python
+
+    import xbatcher
+
+    for batch in da.batch.generator({'time': 10}):
+        pass
+        # actually feed to machine learning library
+    batch
+
 .. toctree::
    :maxdepth: 2
    :caption: Contents:

setup.cfg

@@ -7,7 +7,7 @@ select = B,C,E,F,W,T4,B9
 
 [isort]
 known_first_party=xbatcher
-known_third_party=numpy,pytest,setuptools,xarray
+known_third_party=numpy,pytest,setuptools,sphinx_autosummary_accessors,xarray
 multi_line_output=3
 include_trailing_comma=True
 force_grid_wrap=0

setup.py

@@ -16,8 +16,9 @@
     'Intended Audience :: Science/Research',
     'Programming Language :: Python',
     'Programming Language :: Python :: 3',
-    'Programming Language :: Python :: 3.5',
-    'Programming Language :: Python :: 3.6',
+    'Programming Language :: Python :: 3.7',
+    'Programming Language :: Python :: 3.8',
+    'Programming Language :: Python :: 3.9',
     'Topic :: Scientific/Engineering',
 ]
 

xbatcher/__init__.py

@@ -1 +1,2 @@
+from .accessors import BatchAccessor  # noqa: F401
 from .generators import BatchGenerator  # noqa: F401

xbatcher/accessors.py

@@ -0,0 +1,26 @@
+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):
+        '''
+        Batch accessor returning a BatchGenerator object via the `generator method`
+        '''
+        self._obj = xarray_obj
+
+    def generator(self, *args, **kwargs):
+        '''
+        Return a BatchGenerator via the batch accessor
+
+        Parameters
+        ----------
+        *args : iterable
+            Positional arguments to pass to the `BatchGenerator` constructor.
+        **kwargs : dict
+            Keyword arguments to pass to the `BatchGenerator` constructor.
+        '''
+        return BatchGenerator(self._obj, *args, **kwargs)

xbatcher/tests/test_accessors.py

@@ -0,0 +1,40 @@
+import numpy as np
+import pytest
+import xarray as xr
+
+import xbatcher  # noqa: F401
+from xbatcher import BatchGenerator
+
+
+@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)