Session.virtualfile_in: Refactor the 'check_kind' parameter

pygmt/clib/session.py

@@ -1757,6 +1757,7 @@ def virtualfile_from_stringio(
     def virtualfile_in(
         self,
         check_kind=None,
+        kind=None,
         data=None,
         x=None,
         y=None,
@@ -1815,7 +1816,9 @@ def virtualfile_in(
         ...             print(fout.read().strip())
         <vector memory>: N = 3 <7/9> <4/6> <1/3>
         """
-        kind = data_kind(data, required=required)
+        # Determine the data kind if not given.
+        if kind is None:
+            kind = data_kind(data, required=required, check_kind=check_kind)
         _validate_data_input(
             data=data,
             x=x,
@@ -1826,18 +1829,6 @@ def virtualfile_in(
             kind=kind,
         )
 
-        if check_kind:
-            valid_kinds = ("file", "arg") if required is False else ("file",)
-            if check_kind == "raster":
-                valid_kinds += ("grid", "image")
-            elif check_kind == "vector":
-                valid_kinds += ("empty", "matrix", "vectors", "geojson")
-            if kind not in valid_kinds:
-                raise GMTTypeError(
-                    type(data),
-                    reason=f"Unrecognized data type for {check_kind!r} kind.",
-                )
-
         # Decide which virtualfile_from_ function to use
         _virtualfile_from = {
             "arg": contextlib.nullcontext,

pygmt/helpers/utils.py

@@ -19,7 +19,7 @@
 import xarray as xr
 from pygmt._typing import PathLike
 from pygmt.encodings import charset
-from pygmt.exceptions import GMTInvalidInput, GMTValueError
+from pygmt.exceptions import GMTInvalidInput, GMTTypeError, GMTValueError
 
 # Type hints for the list of encodings supported by PyGMT.
 Encoding = Literal[
@@ -42,6 +42,11 @@
     "ISO-8859-16",
 ]
 
+# Type hints for the list of data kinds.
+Kind = Literal[
+    "arg", "empty", "file", "geojson", "grid", "image", "matrix", "stringio", "vectors"
+]
+
 
 def _validate_data_input(  # ruff: ignore[too-many-branches]
     data=None, x=None, y=None, z=None, required=True, mincols=2, kind=None
@@ -273,11 +278,11 @@ def _check_encoding(argstr: str) -> Encoding:
     return "ISOLatin1+"
 
 
-def data_kind(
-    data: Any, required: bool = True
-) -> Literal[
-    "arg", "empty", "file", "geojson", "grid", "image", "matrix", "stringio", "vectors"
-]:
+def data_kind(  # ruff: ignore[too-many-branches]
+    data: Any,
+    required: bool = True,
+    check_kind: Kind | Sequence[Kind] | Literal["raster", "vector"] | None = None,
+) -> Kind:
     r"""
     Check the kind of data that is provided to a module.
 
@@ -308,6 +313,14 @@ def data_kind(
     required
         Whether 'data' is required. Set to ``False`` when dealing with optional virtual
         files.
+    check_kind
+        Used to validate the type of data that can be passed in. Valid values are:
+
+        - Any recognized data kind
+        - A list/tuple of recognized data kinds
+        - ``"raster"``: shorthand for a sequence of raster-like data kinds
+        - ``"vector"``: shorthand for a sequence of vector-like data kinds
+        - ``None``: means no validatation.
 
     Returns
     -------
@@ -415,6 +428,22 @@ def data_kind(
             kind = "matrix"
         case _:  # Fall back to "vectors" if data is None and required=True.
             kind = "vectors"
+
+    # Now start to check if the data kind is valid.
+    if check_kind is not None:
+        valid_kinds = ("file", "arg") if required is False else ("file",)
+        match check_kind:
+            case "raster":
+                valid_kinds += ("grid", "image")
+            case "vector":
+                valid_kinds += ("empty", "matrix", "vectors", "geojson")
+            case str():
+                valid_kinds = (check_kind,)
+            case list() | tuple():
+                valid_kinds = check_kind
+
+        if kind not in valid_kinds:
+            raise GMTTypeError(dtype=type(data))
     return kind  # type: ignore[return-value]
 
 

pygmt/src/grdcut.py

@@ -115,7 +115,7 @@ def grdcut(
         raise GMTValueError(kind, description="raster kind", choices=["grid", "image"])
 
     # Determine the output data kind based on the input data kind.
-    match inkind := data_kind(grid):
+    match inkind := data_kind(grid, check_kind="raster"):
         case "grid" | "image":
             outkind = inkind
         case "file":
@@ -132,7 +132,7 @@ def grdcut(
 
     with Session() as lib:
         with (
-            lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,
+            lib.virtualfile_in(data=grid, kind=inkind) as vingrd,
             lib.virtualfile_out(kind=outkind, fname=outgrid) as voutgrd,
         ):
             aliasdict["G"] = voutgrd

pygmt/src/legend.py

@@ -131,9 +131,7 @@ def legend(
     if height is not None and width is None:
         width = 0
 
-    kind = data_kind(spec)
-    if kind not in {"empty", "file", "stringio"}:
-        raise GMTTypeError(type(spec))
+    kind = data_kind(spec, check_kind=("empty", "file", "stringio"))
     if kind == "file" and is_nonstr_iter(spec):
         raise GMTTypeError(
             type(spec), reason="Only one legend specification file is allowed."
@@ -161,7 +159,7 @@ def legend(
 
     self._activate_figure()
     with Session() as lib:
-        with lib.virtualfile_in(data=spec, required=False) as vintbl:
+        with lib.virtualfile_in(data=spec, required=False, kind=kind) as vintbl:
             lib.call_module(
                 module="legend", args=build_arg_list(aliasdict, infile=vintbl)
             )

pygmt/src/meca.py

@@ -49,7 +49,7 @@ def _preprocess_spec(spec, colnames, override_cols):
         Dictionary of column names and values to override in the input data. Only makes
         sense if ``spec`` is a dict or :class:`pandas.DataFrame`.
     """
-    kind = data_kind(spec)  # Determine the kind of the input data.
+    kind = data_kind(spec, check_kind="vector")  # Determine the kind of the input data.
 
     # Convert pandas.DataFrame and numpy.ndarray to dict.
     if isinstance(spec, pd.DataFrame):
@@ -388,7 +388,7 @@ def meca(
 
     self._activate_figure()
     with Session() as lib:
-        with lib.virtualfile_in(check_kind="vector", data=spec) as vintbl:
+        with lib.virtualfile_in(data=spec) as vintbl:
             lib.call_module(
                 module="meca", args=build_arg_list(aliasdict, infile=vintbl)
             )

pygmt/src/plot.py

@@ -251,8 +251,9 @@ def plot(  # ruff: ignore[too-many-branches]
     """
     # TODO(GMT>6.5.0): Remove the note for the upstream bug of the "straight_line"
     # parameter.
-    kind = data_kind(data)
+    kind = data_kind(data, check_kind="vector")
     if kind == "empty":  # Data is given via a series of vectors.
+        kind = "vectors"
         data = {"x": x, "y": y}
         # Parameters for vector styles
         if (
@@ -319,7 +320,7 @@ def plot(  # ruff: ignore[too-many-branches]
 
     self._activate_figure()
     with Session() as lib:
-        with lib.virtualfile_in(check_kind="vector", data=data) as vintbl:
+        with lib.virtualfile_in(data=data, kind=kind) as vintbl:
             lib.call_module(
                 module="plot", args=build_arg_list(aliasdict, infile=vintbl)
             )

pygmt/src/plot3d.py

@@ -223,6 +223,7 @@ def plot3d(  # ruff: ignore[too-many-branches]
     # parameter.
     kind = data_kind(data)
     if kind == "empty":  # Data is given via a series of vectors.
+        kind = "vectors"
         data = {"x": x, "y": y, "z": z}
         # Parameters for vector styles
         if (
@@ -292,7 +293,7 @@ def plot3d(  # ruff: ignore[too-many-branches]
 
     self._activate_figure()
     with Session() as lib:
-        with lib.virtualfile_in(check_kind="vector", data=data, mincols=3) as vintbl:
+        with lib.virtualfile_in(data=data, mincols=3, kind=kind) as vintbl:
             lib.call_module(
                 module="plot3d", args=build_arg_list(aliasdict, infile=vintbl)
             )

pygmt/src/text.py

@@ -202,7 +202,7 @@ def text(  # ruff: ignore[too-many-branches, too-many-statements]
         raise GMTParameterError(at_most_one=["textfiles", "position/text", "x/y/text"])
 
     data_is_required = position is None
-    kind = data_kind(textfiles, required=data_is_required)
+    kind = data_kind(textfiles, required=data_is_required, check_kind="vector")
 
     if position is not None:
         if text is None:
@@ -246,6 +246,7 @@ def text(  # ruff: ignore[too-many-branches, too-many-statements]
     confdict = {}
     data = None
     if kind == "empty":
+        kind = "vectors"
         data = {"x": x, "y": y}
 
         for arg, flag, name in array_args:
@@ -301,7 +302,9 @@ def text(  # ruff: ignore[too-many-branches, too-many-statements]
     self._activate_figure()
     with Session() as lib:
         with lib.virtualfile_in(
-            check_kind="vector", data=textfiles or data, required=data_is_required
+            data=textfiles or data,
+            required=data_is_required,
+            kind=kind,
         ) as vintbl:
             lib.call_module(
                 module="text",

pygmt/src/x2sys_cross.py

@@ -200,7 +200,7 @@ def x2sys_cross(
 
     file_contexts: list[contextlib.AbstractContextManager[Any]] = []
     for track in tracks:
-        match data_kind(track):
+        match data_kind(track, check_kind="vector"):
             case "file":
                 file_contexts.append(contextlib.nullcontext(track))
             case "vectors":