Merge pull request #118 from raybellwaves/transform

sel transform

Author: martindurant

intake_xarray/derived.py

@@ -0,0 +1,64 @@
+from intake import Schema
+from intake.source.derived import GenericTransform
+
+
+class XArrayTransform(GenericTransform):
+    """Transform where the input and output are both xarray objects.
+    You must supply ``transform`` and any ``transform_kwargs``.
+    """
+
+    input_container = "xarray"
+    container = "xarray"
+    optional_params = {}
+    _ds = None
+
+    def to_dask(self):
+        if self._ds is None:
+            self._pick()
+            self._ds = self._transform(
+                self._source.to_dask(), **self._params["transform_kwargs"]
+            )
+        return self._ds
+
+    def _get_schema(self):
+        """load metadata only if needed"""
+        self.to_dask()
+        return Schema(
+            datashape=None,
+            dtype=None,
+            shape=None,
+            npartitions=None,
+            extra_metadata=self._ds.extra_metadata,
+        )
+
+    def read(self):
+        return self.to_dask().compute()
+
+
+class Sel(XArrayTransform):
+    """Simple array transform to subsample an xarray object using
+    the sel method.
+    Note that you could use XArrayTransform directly, by writing a
+    function to choose the subsample instead of a method as here.
+    """
+
+    input_container = "xarray"
+    container = "xarray"
+    required_params = ["indexers"]
+
+    def __init__(self, indexers, **kwargs):
+        """
+        indexers: dict (stord as str) which is passed to xarray.Dataset.sel
+        """
+        # this class wants required "indexers", but XArrayTransform
+        # uses "transform_kwargs", which we don't need since we use a method for the
+        # transform
+        kwargs.update(
+            transform=self.sel,
+            indexers=indexers,
+            transform_kwargs={},
+        )
+        super().__init__(**kwargs)
+
+    def sel(self, ds):
+        return ds.sel(eval(self._params["indexers"]))

intake_xarray/tests/data/catalog.yaml

@@ -114,3 +114,12 @@ sources:
       chunks: {}
       auth: null
       engine: netcdf4
+  xarray_source_sel:
+    description: select subsample of xarray_source entry
+    driver: intake_xarray.derived.XArrayTransform
+    args:
+      targets:
+        - xarray_source
+      transform: "intake_xarray.tests.test_derived._sel"
+      transform_kwargs:
+        indexers: "dict([('lat', 20)])"

intake_xarray/tests/test_derived.py

@@ -0,0 +1,23 @@
+import os
+import pytest
+
+from intake import open_catalog
+from xarray.tests import assert_allclose
+
+
+# Function used in xarray_source_sel entry in catalog.yaml
+def _sel(ds, indexers: str):
+    """indexers: dict (stored as str) which is passed to xarray.Dataset.sel"""
+    return ds.sel(eval(indexers))
+
+
+@pytest.fixture
+def catalog():
+    path = os.path.dirname(__file__)
+    return open_catalog(os.path.join(path, "data", "catalog.yaml"))
+
+
+def test_catalog(catalog):
+    expected = catalog["xarray_source"].read().sel(lat=20)
+    actual = catalog["xarray_source_sel"].read()
+    assert_allclose(actual, expected)