Add very minimal xarray plugin for engine="rasterio" (#281)

Author: alexamici

docs/getting_started/getting_started.rst

@@ -26,6 +26,12 @@ Why use :func:`rioxarray.open_rasterio` instead of `xarray.open_rasterio`?
 5. It adds the coordinate axis CF metadata.
 6. It loads raster metadata into the attributes.
 
+rioxarray 0.4+ enables passing `engine="rasterio"` to ``xarray.open_dataset`` and ``xarray.open_mfdataset`` for xarray 0.18+:
+
+.. code-block:: python
+
+    ds = xr.open_dataset("my.tif", engine="rasterio")
+
 
 .. toctree::
    :maxdepth: 1

docs/history.rst

@@ -6,6 +6,7 @@ Latest
 - DEP: Python 3.6+ (issue #215)
 - DEP: xarray 0.17+ (needed for issue #282)
 - REF: Store `grid_mapping` in `encoding` instead of `attrs` (issue #282)
+- ENH: enable `engine="rasterio"` via xarray backend API (issue #197 pull #281)
 
 0.3.2
 -----

rioxarray/xarray_plugin.py

@@ -0,0 +1,63 @@
+import os.path
+
+import xarray as xr
+
+from . import _io
+
+CAN_OPEN_EXTS = {
+    "asc",
+    "geotif",
+    "geotiff",
+    "img",
+    "j2k",
+    "jp2",
+    "jpg",
+    "jpeg",
+    "png",
+    "tif",
+    "tiff",
+    "vrt",
+}
+
+
+class RasterioBackend(xr.backends.common.BackendEntrypoint):
+    """
+    .. versionadded:: 0.4
+    """
+
+    def open_dataset(
+        self,
+        filename_or_obj,
+        drop_variables=None,
+        mask_and_scale=True,
+        parse_coordinates=None,
+        lock=None,
+        masked=False,
+        variable=None,
+        group=None,
+        default_name="band_data",
+        open_kwargs=None,
+    ):
+        if open_kwargs is None:
+            open_kwargs = {}
+        ds = _io.open_rasterio(
+            filename_or_obj,
+            mask_and_scale=mask_and_scale,
+            parse_coordinates=parse_coordinates,
+            lock=lock,
+            masked=masked,
+            variable=variable,
+            group=group,
+            default_name=default_name,
+            **open_kwargs,
+        )
+        if isinstance(ds, xr.DataArray):
+            ds = ds.to_dataset()
+        return ds
+
+    def guess_can_open(self, filename_or_obj):
+        try:
+            _, ext = os.path.splitext(filename_or_obj)
+        except TypeError:
+            return False
+        return ext[1:].lower() in CAN_OPEN_EXTS

setup.py

@@ -71,4 +71,7 @@ def get_version():
     version=get_version(),
     zip_safe=False,
     python_requires=">=3.7",
+    entry_points={
+        "xarray.backends": ["rasterio=rioxarray.xarray_plugin:RasterioBackend"]
+    },
 )

test/integration/test_integration_xarray_plugin.py

@@ -0,0 +1,26 @@
+import os.path
+
+import pytest
+
+from test.conftest import TEST_INPUT_DATA_DIR
+
+# FIXME: change to the next xarray version after release
+xr = pytest.importorskip("xarray", minversion="0.17.1.dev0")
+
+
+def test_xarray_open_dataset():
+    cog_file = os.path.join(TEST_INPUT_DATA_DIR, "cog.tif")
+
+    ds = xr.open_dataset(cog_file, engine="rasterio")
+
+    assert isinstance(ds, xr.Dataset)
+    assert "band_data" in ds.data_vars
+    assert ds.data_vars["band_data"].shape == (1, 500, 500)
+    assert "spatial_ref" not in ds.data_vars
+    assert "spatial_ref" in ds.coords
+    assert "grid_mapping" not in ds.data_vars["band_data"].attrs
+    assert "grid_mapping" in ds.data_vars["band_data"].encoding
+
+    ds = xr.open_dataset(cog_file)
+
+    assert isinstance(ds, xr.Dataset)