diff --git a/docs/getting_started/getting_started.rst b/docs/getting_started/getting_started.rst index fe8241d3..e47c4725 100644 --- a/docs/getting_started/getting_started.rst +++ b/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 diff --git a/docs/history.rst b/docs/history.rst index 38575cf1..7c926b49 100644 --- a/docs/history.rst +++ b/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 ----- diff --git a/rioxarray/xarray_plugin.py b/rioxarray/xarray_plugin.py new file mode 100644 index 00000000..c7787fa3 --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index 69f24ce9..09a09c18 100644 --- a/setup.py +++ b/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"] + }, ) diff --git a/test/integration/test_integration_xarray_plugin.py b/test/integration/test_integration_xarray_plugin.py new file mode 100644 index 00000000..2fbb06a8 --- /dev/null +++ b/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)