Allow passing RGB xarray.DataArray images into grdimage

pygmt/src/grdimage.py

@@ -5,6 +5,7 @@
 
 from pygmt.clib import Session
 from pygmt.helpers import (
+    GMTTempFile,
     build_arg_string,
     data_kind,
     fmt_docstring,
@@ -179,17 +180,25 @@ def grdimage(self, grid, **kwargs):
     >>> fig.show()
     """
     kwargs = self._preprocess(**kwargs)  # pylint: disable=protected-access
-    with Session() as lib:
-        file_context = lib.virtualfile_from_data(check_kind="raster", data=grid)
-        with contextlib.ExitStack() as stack:
-            # shading using an xr.DataArray
-            if kwargs.get("I") is not None and data_kind(kwargs["I"]) == "grid":
-                shading_context = lib.virtualfile_from_data(
-                    check_kind="raster", data=kwargs["I"]
-                )
-                kwargs["I"] = stack.enter_context(shading_context)
 
-            fname = stack.enter_context(file_context)
-            lib.call_module(
-                module="grdimage", args=build_arg_string(kwargs, infile=fname)
-            )
+    with GMTTempFile(suffix=".tif") as tmpfile:
+        if hasattr(grid, "dims") and len(grid.dims) == 3:
+            grid.rio.to_raster(raster_path=tmpfile.name)
+            _grid = tmpfile.name
+        else:
+            _grid = grid
+
+        with Session() as lib:
+            file_context = lib.virtualfile_from_data(check_kind="raster", data=_grid)
+            with contextlib.ExitStack() as stack:
+                # shading using an xr.DataArray
+                if kwargs.get("I") is not None and data_kind(kwargs["I"]) == "grid":
+                    shading_context = lib.virtualfile_from_data(
+                        check_kind="raster", data=kwargs["I"]
+                    )
+                    kwargs["I"] = stack.enter_context(shading_context)
+
+                fname = stack.enter_context(file_context)
+                lib.call_module(
+                    module="grdimage", args=build_arg_string(kwargs, infile=fname)
+                )

pygmt/tests/baseline/test_grdimage_image.png.dvc

@@ -0,0 +1,4 @@
+outs:
+- md5: 2e919645d5af956ec4f8aa054a86a70a
+  size: 110214
+  path: test_grdimage_image.png

pygmt/tests/test_grdimage_image.py

@@ -0,0 +1,63 @@
+"""
+Test Figure.grdimage on 3-band RGB images.
+"""
+import numpy as np
+import pandas as pd
+import pytest
+import xarray as xr
+from pygmt import Figure, which
+
+rasterio = pytest.importorskip("rasterio")
+rioxarray = pytest.importorskip("rioxarray")
+
+
+@pytest.fixture(scope="module", name="xr_image")
+def fixture_xr_image():
+    """
+    Load the image data from Blue Marble as an xarray.DataArray with shape
+    {"band": 3, "y": 180, "x": 360}.
+    """
+    geotiff = which(fname="@earth_day_01d_p", download="c")
+    with rioxarray.open_rasterio(filename=geotiff) as rda:
+        if len(rda.band) == 1:
+            with rasterio.open(fp=geotiff) as src:
+                df_colormap = pd.DataFrame.from_dict(
+                    data=src.colormap(1), orient="index"
+                )
+                array = src.read()
+
+                red = np.vectorize(df_colormap[0].get)(array)
+                green = np.vectorize(df_colormap[1].get)(array)
+                blue = np.vectorize(df_colormap[2].get)(array)
+                # alpha = np.vectorize(df_colormap[3].get)(array)
+
+            rda.data = red
+            da_red = rda.astype(dtype=np.uint8).copy()
+            rda.data = green
+            da_green = rda.astype(dtype=np.uint8).copy()
+            rda.data = blue
+            da_blue = rda.astype(dtype=np.uint8).copy()
+
+            xr_image = xr.concat(objs=[da_red, da_green, da_blue], dim="band")
+        assert xr_image.sizes == {"band": 3, "y": 180, "x": 360}
+        return xr_image
+
+
+@pytest.mark.mpl_image_compare
+def test_grdimage_image():
+    """
+    Plot a 3-band RGB image using file input.
+    """
+    fig = Figure()
+    fig.grdimage(grid="@earth_day_01d")
+    return fig
+
+
+@pytest.mark.mpl_image_compare(filename="test_grdimage_image.png")
+def test_grdimage_image_dataarray(xr_image):
+    """
+    Plot a 3-band RGB image using xarray.DataArray input.
+    """
+    fig = Figure()
+    fig.grdimage(grid=xr_image)
+    return fig