Add very minimal xarray plugin for engine="rasterio"

.github/workflows/tests.yaml

@@ -108,8 +108,8 @@ jobs:
           #   pandas \
           #   scipy;
           python -m pip install --upgrade \
-            git+https://github.com/dask/dask.git@master \
-            git+https://github.com/dask/distributed.git@master \
+            git+https://github.com/dask/dask.git@main \
+            git+https://github.com/dask/distributed.git@main \
             git+https://github.com/mapbox/rasterio.git@master \
             git+https://github.com/pyproj4/pyproj.git@master \
             git+https://github.com/pydata/xarray.git@master;

rioxarray/xarray_plugin.py

@@ -0,0 +1,20 @@
+import os.path
+
+import xarray as xr
+
+from . import _io
+
+
+class GdalBackend(xr.backends.common.BackendEntrypoint):
+    def open_dataset(self, filename_or_obj, drop_variables=None):
+        ds = _io.open_rasterio(filename_or_obj).to_dataset("band")
+        ds = ds.rename({idx: f"band{idx}" for idx in ds.data_vars})
+        ds = ds.reset_coords("spatial_ref")
+        return ds
+
+    def guess_can_open(self, filename_or_obj):
+        try:
+            _, ext = os.path.splitext(filename_or_obj)
+        except TypeError:
+            return False
+        return ext in {".tif", ".geotif", ".tiff", ".geotiff"}

setup.py

@@ -72,4 +72,5 @@ def get_version():
     version=get_version(),
     zip_safe=False,
     python_requires=">=3.6",
+    entry_points={"xarray.backends": ["gdal=rioxarray.xarray_plugin:GdalBackend"]},
 )

test/integration/test_integration_xarray_plugin.py

@@ -0,0 +1,21 @@
+import os.path
+
+import pytest
+
+from test.conftest import TEST_INPUT_DATA_DIR
+
+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="gdal")
+
+    assert isinstance(ds, xr.Dataset)
+    assert "band1" in ds.data_vars
+    assert ds.data_vars["band1"].shape == (500, 500)
+
+    ds = xr.open_dataset(cog_file)
+
+    assert isinstance(ds, xr.Dataset)