Implement gmt xarray BackendEntrypoint

doc/api/index.rst

@@ -173,6 +173,7 @@ Input/output
 .. autosummary::
     :toctree: generated
 
+    GMTBackendEntrypoint
     load_dataarray
 
 GMT Defaults

pygmt/__init__.py

@@ -65,6 +65,7 @@
     x2sys_init,
     xyz2grd,
 )
+from pygmt.xarray_backend import GMTBackendEntrypoint
 
 # Start our global modern mode session
 _begin()

pygmt/datasets/samples.py

@@ -204,7 +204,7 @@ def _load_earth_relief_holes() -> xr.DataArray:
         is in meters.
     """
     fname = which("@earth_relief_20m_holes.grd", download="c")
-    return load_dataarray(fname, engine="netcdf4")
+    return load_dataarray(fname, engine="gmt", decode_kind="grid")
 
 
 class GMTSampleData(NamedTuple):

pygmt/helpers/testing.py

@@ -154,7 +154,7 @@ def load_static_earth_relief():
         A grid of Earth relief for internal tests.
     """
     fname = which("@static_earth_relief.nc", download="c")
-    return load_dataarray(fname)
+    return load_dataarray(fname, decode_kind="grid")
 
 
 def skip_if_no(package):

pygmt/io.py

@@ -2,10 +2,14 @@
 PyGMT input/output (I/O) utilities.
 """
 
+from typing import Literal
+
 import xarray as xr
 
 
-def load_dataarray(filename_or_obj, **kwargs):
+def load_dataarray(
+    filename_or_obj, *, engine="gmt", decode_kind: Literal["grid", "image"], **kwargs
+):
     """
     Open, load into memory, and close a DataArray from a file or file-like object
     containing a single data variable.
@@ -32,6 +36,10 @@ def load_dataarray(filename_or_obj, **kwargs):
     -------
     datarray : xarray.DataArray
         The newly created DataArray.
+    engine : str
+        Engine to use when reading files. Default engine is 'gmt'.
+    decode_kind
+        The kind of data to read. Valid values are ``"grid"`` or ``"image"``.
 
     See Also
     --------
@@ -41,7 +49,9 @@ def load_dataarray(filename_or_obj, **kwargs):
         msg = "'cache' has no effect in this context."
         raise TypeError(msg)
 
-    with xr.open_dataarray(filename_or_obj, **kwargs) as dataarray:
+    with xr.open_dataarray(
+        filename_or_obj, engine=engine, decode_kind=decode_kind, **kwargs
+    ) as dataarray:
         result = dataarray.load()
         _ = result.gmt  # load GMTDataArray accessor information
 

pygmt/tests/test_xarray_backend.py

@@ -0,0 +1,57 @@
+"""
+Tests for xarray 'gmt' backend engine.
+"""
+
+import re
+
+import pytest
+import xarray as xr
+from pygmt.enums import GridRegistration, GridType
+from pygmt.exceptions import GMTInvalidInput
+
+
+def test_xarray_backend_gmt_read_grid():
+    """
+    Ensure that passing engine='gmt' to xarray.open_dataarray works for reading
+    NetCDF grids.
+    """
+    with xr.open_dataarray(
+        filename_or_obj="@static_earth_relief.nc", engine="gmt", decode_kind="grid"
+    ) as da:
+        assert da.sizes == {"lat": 14, "lon": 8}
+        assert da.dtype == "float32"
+        assert da.gmt.registration == GridRegistration.PIXEL
+        assert da.gmt.gtype == GridType.GEOGRAPHIC
+
+
+def test_xarray_backend_gmt_read_image():
+    """
+    Ensure that passing engine='gmt' to xarray.open_dataarray works for reading
+    GeoTIFF images.
+    """
+    with xr.open_dataarray(
+        filename_or_obj="@earth_day_01d", engine="gmt", decode_kind="image"
+    ) as da:
+        assert da.sizes == {"band": 3, "y": 180, "x": 360}
+        assert da.dtype == "uint8"
+        assert da.gmt.registration == GridRegistration.PIXEL
+        assert da.gmt.gtype == GridType.GEOGRAPHIC
+
+
+def test_xarray_backend_gmt_read_invalid_kind():
+    """
+    Check that xarray.open_dataarray(..., engine="gmt") fails with missing or incorrect
+    'decode_kind'.
+    """
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "GMTBackendEntrypoint.open_dataset() missing 1 required keyword-only argument: 'decode_kind'"
+        ),
+    ):
+        xr.open_dataarray("nokind.nc", engine="gmt")
+
+    with pytest.raises(GMTInvalidInput):
+        xr.open_dataarray(
+            filename_or_obj="invalid.tif", engine="gmt", decode_kind="invalid"
+        )

pygmt/xarray_backend.py

@@ -0,0 +1,72 @@
+"""
+An xarray backend for reading raster grid/image files using the 'gmt' engine.
+"""
+
+import os
+from pathlib import Path
+from typing import Literal
+
+import xarray as xr
+from pygmt.clib import Session
+from pygmt.exceptions import GMTInvalidInput
+from pygmt.helpers import build_arg_list
+from pygmt.src.which import which
+from xarray.backends import BackendEntrypoint
+
+
+class GMTBackendEntrypoint(BackendEntrypoint):
+    """
+    Xarray backend to read raster grid/image files using 'gmt' engine.
+
+    Relies on the libgdal-netcdf driver used by GMT C for NetCDF files, and libgdal for
+    GeoTIFF files.
+    """
+
+    description = "Open raster (.grd, .nc or .tif) files in Xarray via GMT read."
+    open_dataset_parameters = ("filename_or_obj", "kind")
+    url = "https://github.com/GenericMappingTools/pygmt"
+
+    def open_dataset(  # type: ignore[override]
+        self,
+        filename_or_obj: str | os.PathLike,
+        *,
+        drop_variables=None,  # noqa: ARG002
+        decode_kind: Literal["grid", "image"],
+        # other backend specific keyword arguments
+        # `chunks` and `cache` DO NOT go here, they are handled by xarray
+    ) -> xr.Dataset:
+        """
+        Backend open_dataset method used by Xarray in :py:func:`~xarray.open_dataset`.
+        """
+        if decode_kind not in {"grid", "image"}:
+            msg = f"Invalid raster kind: '{decode_kind}'. Valid values are 'grid' or 'image'."
+            raise GMTInvalidInput(msg)
+
+        with Session() as lib:
+            with lib.virtualfile_out(kind=decode_kind) as voutfile:
+                kwdict = {"T": {"grid": "g", "image": "i"}[decode_kind]}
+                lib.call_module(
+                    module="read",
+                    args=[filename_or_obj, voutfile, *build_arg_list(kwdict)],
+                )
+
+                raster: xr.DataArray = lib.virtualfile_to_raster(
+                    vfname=voutfile, kind=decode_kind
+                )
+                # Add "source" encoding
+                source = which(fname=filename_or_obj)
+                raster.encoding["source"] = (
+                    source[0] if isinstance(source, list) else source
+                )
+                _ = raster.gmt  # Load GMTDataArray accessor information
+                return raster.to_dataset()
+
+    def guess_can_open(self, filename_or_obj) -> bool:
+        """
+        Backend open_dataset method used by Xarray in :py:func:`~xarray.open_dataset`.
+        """
+        try:
+            ext = Path(filename_or_obj).suffix
+        except TypeError:
+            return False
+        return ext in {".grd", ".nc", ".tif"}

pyproject.toml

@@ -52,6 +52,9 @@ all = [
     "rioxarray",
 ]
 
+[project.entry-points."xarray.backends"]
+gmt = "pygmt.xarray_backend:GMTBackendEntrypoint"
+
 [project.urls]
 "Homepage" = "https://www.pygmt.org"
 "Documentation" = "https://www.pygmt.org"