diff --git a/pygmt/helpers/utils.py b/pygmt/helpers/utils.py index a7092a7ea9b..286f0430c9c 100644 --- a/pygmt/helpers/utils.py +++ b/pygmt/helpers/utils.py @@ -14,7 +14,74 @@ from pygmt.exceptions import GMTInvalidInput -def data_kind(data, x=None, y=None, z=None, required_z=False): +def _validate_data_input( + data=None, x=None, y=None, z=None, required_z=False, kind=None +): + """ + Check if the combination of data/x/y/z is valid. + + Examples + -------- + >>> _validate_data_input(data="infile") + >>> _validate_data_input(x=[1, 2, 3], y=[4, 5, 6]) + >>> _validate_data_input(x=[1, 2, 3], y=[4, 5, 6], z=[7, 8, 9]) + >>> _validate_data_input() + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: No input data provided. + >>> _validate_data_input(x=[1, 2, 3]) + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Must provide both x and y. + >>> _validate_data_input(y=[4, 5, 6]) + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Must provide both x and y. + >>> _validate_data_input(x=[1, 2, 3], y=[4, 5, 6], required_z=True) + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Must provide x, y, and z. + >>> _validate_data_input(data="infile", x=[1, 2, 3]) + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Too much data. Use either data or x/y/z. + >>> _validate_data_input(data="infile", y=[4, 5, 6]) + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Too much data. Use either data or x/y/z. + >>> _validate_data_input(data="infile", z=[7, 8, 9]) + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Too much data. Use either data or x/y/z. + + Raises + ------ + GMTInvalidInput + If the data input is not valid. + """ + if data is None: # data is None + if x is None and y is None: # both x and y are None + raise GMTInvalidInput("No input data provided.") + if x is None or y is None: # either x or y is None + raise GMTInvalidInput("Must provide both x and y.") + # both x and y are not None, now check z + if required_z and z is None: + raise GMTInvalidInput("Must provide x, y, and z.") + else: # data is not None + if x is not None or y is not None or z is not None: + raise GMTInvalidInput("Too much data. Use either data or x/y/z.") + # For 'matrix' kind, check if data has the required z column + if kind == "matrix" and required_z: + if hasattr(data, "shape"): # np.ndarray or pd.DataFrame + if len(data.shape) == 1 and data.shape[0] < 3: + raise GMTInvalidInput("data must provide x, y, and z columns.") + if len(data.shape) > 1 and data.shape[1] < 3: + raise GMTInvalidInput("data must provide x, y, and z columns.") + if hasattr(data, "data_vars") and len(data.data_vars) < 3: # xr.Dataset + raise GMTInvalidInput("data must provide x, y, and z columns.") + + +def data_kind(data=None, x=None, y=None, z=None, required_z=False): """ Check what kind of data is provided to a module. @@ -64,15 +131,6 @@ def data_kind(data, x=None, y=None, z=None, required_z=False): >>> data_kind(data=xr.DataArray(np.random.rand(4, 3))) 'grid' """ - if data is None and x is None and y is None: - raise GMTInvalidInput("No input data provided.") - if data is not None and (x is not None or y is not None or z is not None): - raise GMTInvalidInput("Too much data. Use either data or x and y.") - if data is None and (x is None or y is None): - raise GMTInvalidInput("Must provide both x and y.") - if data is None and required_z and z is None: - raise GMTInvalidInput("Must provide x, y, and z.") - if isinstance(data, (str, pathlib.PurePath)): kind = "file" elif isinstance(data, xr.DataArray): @@ -80,14 +138,10 @@ def data_kind(data, x=None, y=None, z=None, required_z=False): elif hasattr(data, "__geo_interface__"): kind = "geojson" elif data is not None: - if required_z and ( - getattr(data, "shape", (3, 3))[1] < 3 # np.array, pd.DataFrame - or len(getattr(data, "data_vars", (0, 1, 2))) < 3 # xr.Dataset - ): - raise GMTInvalidInput("data must provide x, y, and z columns.") kind = "matrix" else: kind = "vectors" + _validate_data_input(data=data, x=x, y=y, z=z, required_z=required_z, kind=kind) return kind