Alias drapegrid (G) for grdview

Included one test for a valid (xarray.DataArray) drapegrid input, and one test for an invalid input. Also using contextlib.ExitStack() as we have more than one file_context (one for reliefgrid, one for drapegrid) to enter and exit properly (see https://stackoverflow.com/questions/3024925/create-a-with-block-on-several-context-managers).

Author: weiji14

pygmt/base_plotting.py

@@ -2,6 +2,8 @@
 Base class with plot generating commands.
 Does not define any special non-GMT methods (savefig, show, etc).
 """
+import contextlib
+
 from .clib import Session
 from .exceptions import GMTInvalidInput
 from .helpers import (
@@ -240,6 +242,7 @@ def grdimage(self, grid, **kwargs):
         JZ="zsize",
         B="frame",
         C="cmap",
+        G="drapegrid",
         N="plane",
         Q="surftype",
         Wc="contourpen",
@@ -271,6 +274,13 @@ def grdview(self, reliefgrid, **kwargs):
         cmap (C) : str
             The name of the color palette table to use.
 
+        drapegrid (G) : str or xarray.DataArray
+            The file name or a DataArray of the image grid to be draped on top of the
+            relief provided by reliefgrid. [Default determines colors from reliefgrid].
+            Note that -Jz and -N always refers to the reliefgrid. The drapegrid only
+            provides the information pertaining to colors, which (if drapegrid is a
+            grid) will be looked-up via the CPT (see -C).
+
         plane (N) : float or str
             ``level[+gfill]``.
             Draws a plane at this z-level. If the optional color is provided via the +g
@@ -312,9 +322,22 @@ def grdview(self, reliefgrid, **kwargs):
                 file_context = lib.virtualfile_from_grid(reliefgrid)
             else:
                 raise GMTInvalidInput(
-                    "Unrecognized data type: {}".format(type(reliefgrid))
+                    f"Unrecognized data type for reliefgrid: {type(reliefgrid)}"
                 )
-            with file_context as fname:
+
+            with contextlib.ExitStack() as stack:
+                fname = stack.enter_context(file_context)
+                if "G" in kwargs:
+                    drapegrid = kwargs["G"]
+                    if data_kind(drapegrid) in ("file", "grid"):
+                        if data_kind(drapegrid) == "grid":
+                            drape_context = lib.virtualfile_from_grid(drapegrid)
+                            drapefile = stack.enter_context(drape_context)
+                            kwargs["G"] = drapefile
+                    else:
+                        raise GMTInvalidInput(
+                            f"Unrecognized data type for drapegrid: {type(drapegrid)}"
+                        )
                 arg_str = " ".join([fname, build_arg_string(kwargs)])
                 lib.call_module("grdview", arg_str)
 

pygmt/tests/baseline/test_grdview_drapefile_dataarray.png

@@ -200,3 +200,28 @@ def test_grdview_on_a_plane_styled_with_facadepen(grid):
         facadepen="0.5p,blue,dash",
     )
     return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_grdview_drapegrid_dataarray(grid):
+    """
+    Run grdview by passing in both a reliefgrid and drapegrid as an xarray.DataArray,
+    setting a colormap for producing an image plot.
+    """
+    drapegrid = 1.1 * grid
+
+    fig = Figure()
+    fig.grdview(reliefgrid=grid, drapegrid=drapegrid, cmap="oleron", surftype="c")
+    return fig
+
+
+def test_grdview_wrong_kind_of_drapegrid(grid):
+    """
+    Run grdview using drapegrid input that is not an xarray.DataArray or file.
+    """
+    dataset = grid.to_dataset()  # convert xarray.DataArray to xarray.Dataset
+    assert data_kind(dataset) == "matrix"
+
+    fig = Figure()
+    with pytest.raises(GMTInvalidInput):
+        fig.grdview(reliefgrid=grid, drapegrid=dataset)