diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index dda4b03ae5c..c9479106d46 100644 --- a/pygmt/clib/session.py +++ b/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()) : 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, diff --git a/pygmt/helpers/utils.py b/pygmt/helpers/utils.py index b16885689ff..5c595ba22ea 100644 --- a/pygmt/helpers/utils.py +++ b/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] diff --git a/pygmt/src/grdcut.py b/pygmt/src/grdcut.py index a8ed5ea6eb0..a715e51f03a 100644 --- a/pygmt/src/grdcut.py +++ b/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 diff --git a/pygmt/src/legend.py b/pygmt/src/legend.py index 5e5f320e3ac..af9db9beef8 100644 --- a/pygmt/src/legend.py +++ b/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) ) diff --git a/pygmt/src/meca.py b/pygmt/src/meca.py index a2c84cd641a..0313d78d571 100644 --- a/pygmt/src/meca.py +++ b/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) ) diff --git a/pygmt/src/plot.py b/pygmt/src/plot.py index bb2d5097e76..6d689a46651 100644 --- a/pygmt/src/plot.py +++ b/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) ) diff --git a/pygmt/src/plot3d.py b/pygmt/src/plot3d.py index 00aae4b1e18..c1240bd2797 100644 --- a/pygmt/src/plot3d.py +++ b/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) ) diff --git a/pygmt/src/text.py b/pygmt/src/text.py index 4e9bbf5177d..5e80e139d6a 100644 --- a/pygmt/src/text.py +++ b/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", diff --git a/pygmt/src/x2sys_cross.py b/pygmt/src/x2sys_cross.py index 174790ebe40..26bec531a15 100644 --- a/pygmt/src/x2sys_cross.py +++ b/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":