pygmt.grdcut: Support both grid and image output

Author: seisman

pygmt/clib/session.py

@@ -25,14 +25,15 @@
     vectors_to_arrays,
 )
 from pygmt.clib.loading import load_libgmt
-from pygmt.datatypes import _GMT_DATASET, _GMT_GRID
+from pygmt.datatypes import _GMT_DATASET, _GMT_GRID, _GMT_IMAGE
 from pygmt.exceptions import (
     GMTCLibError,
     GMTCLibNoSessionError,
     GMTInvalidInput,
     GMTVersionError,
 )
 from pygmt.helpers import (
+    GMTTempFile,
     data_kind,
     fmt_docstring,
     tempfile_from_geojson,
@@ -1649,7 +1650,9 @@ def virtualfile_in(  # noqa: PLR0912
 
     @contextlib.contextmanager
     def virtualfile_out(
-        self, kind: Literal["dataset", "grid"] = "dataset", fname: str | None = None
+        self,
+        kind: Literal["dataset", "grid", "image"] = "dataset",
+        fname: str | None = None,
     ):
         r"""
         Create a virtual file or an actual file for storing output data.
@@ -1706,8 +1709,11 @@ def virtualfile_out(
             family, geometry = {
                 "dataset": ("GMT_IS_DATASET", "GMT_IS_PLP"),
                 "grid": ("GMT_IS_GRID", "GMT_IS_SURFACE"),
+                "image": ("GMT_IS_IMAGE", "GMT_IS_SURFACE"),
             }[kind]
-            with self.open_virtualfile(family, geometry, "GMT_OUT", None) as vfile:
+            with self.open_virtualfile(
+                family, geometry, "GMT_OUT|GMT_IS_REFERENCE", None
+            ) as vfile:
                 yield vfile
 
     def inquire_virtualfile(self, vfname: str) -> int:
@@ -1801,9 +1807,13 @@ def read_virtualfile(
         # _GMT_DATASET).
         if kind is None:  # Return the ctypes void pointer
             return pointer
-        if kind in ["image", "cube"]:
+        if kind == "cube":
             raise NotImplementedError(f"kind={kind} is not supported yet.")
-        dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID}[kind]
+        dtype = {
+            "dataset": _GMT_DATASET,
+            "grid": _GMT_GRID,
+            "image": _GMT_IMAGE,
+        }[kind]
         return ctp.cast(pointer, ctp.POINTER(dtype))
 
     def virtualfile_to_dataset(
@@ -2013,6 +2023,10 @@ def virtualfile_to_raster(
                 self["GMT_IS_IMAGE"]: "image",
                 self["GMT_IS_CUBE"]: "cube",
             }[family]
+        if kind == "image":
+            with GMTTempFile(suffix=".tif") as tmpfile:
+                self.call_module("write", f"{vfname} {tmpfile.name} -Ti")
+                return xr.load_dataarray(tmpfile.name)
         return self.read_virtualfile(vfname, kind=kind).contents.to_dataarray()
 
     def extract_region(self):

pygmt/datatypes/__init__.py

@@ -4,3 +4,4 @@
 
 from pygmt.datatypes.dataset import _GMT_DATASET
 from pygmt.datatypes.grid import _GMT_GRID
+from pygmt.datatypes.image import _GMT_IMAGE

pygmt/datatypes/image.py

@@ -0,0 +1,9 @@
+"""
+Wrapper for the GMT_GRID data type.
+"""
+
+import ctypes as ctp
+
+
+class _GMT_IMAGE(ctp.Structure):  # noqa: N801
+    pass

pygmt/src/grdcut.py

@@ -4,20 +4,19 @@
 
 from pygmt.clib import Session
 from pygmt.helpers import (
-    GMTTempFile,
     build_arg_string,
+    data_kind,
     fmt_docstring,
     kwargs_to_strings,
     use_alias,
 )
-from pygmt.io import load_dataarray
+from pygmt.src.which import which
 
 __doctest_skip__ = ["grdcut"]
 
 
 @fmt_docstring
 @use_alias(
-    G="outgrid",
     R="region",
     J="projection",
     N="extend",
@@ -27,7 +26,7 @@
     f="coltypes",
 )
 @kwargs_to_strings(R="sequence")
-def grdcut(grid, **kwargs):
+def grdcut(grid, outgrid: str | None = None, **kwargs):
     r"""
     Extract subregion from a grid.
 
@@ -99,13 +98,24 @@ def grdcut(grid, **kwargs):
     >>> # 12° E to 15° E and a latitude range of 21° N to 24° N
     >>> new_grid = pygmt.grdcut(grid=grid, region=[12, 15, 21, 24])
     """
-    with GMTTempFile(suffix=".nc") as tmpfile:
-        with Session() as lib:
-            with lib.virtualfile_in(check_kind="raster", data=grid) as vingrd:
-                if (outgrid := kwargs.get("G")) is None:
-                    kwargs["G"] = outgrid = tmpfile.name  # output to tmpfile
-                lib.call_module(
-                    module="grdcut", args=build_arg_string(kwargs, infile=vingrd)
-                )
+    inkind = data_kind(grid)
+    outkind = "image" if inkind == "image" else "grid"
+    if inkind == "file":
+        realpath = which(str(grid))
+        if isinstance(realpath, list):
+            realpath = realpath[0]
+        if realpath.endswith(".tif"):
+            outkind = "image"
 
-        return load_dataarray(outgrid) if outgrid == tmpfile.name else None
+    with Session() as lib:
+        with (
+            lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,
+            lib.virtualfile_out(kind=outkind, fname=outgrid) as voutgrd,
+        ):
+            kwargs["G"] = voutgrd
+            lib.call_module(
+                module="grdcut", args=build_arg_string(kwargs, infile=vingrd)
+            )
+            return lib.virtualfile_to_raster(
+                outgrid=outgrid, kind=outkind, vfname=voutgrd
+            )