Raise GMTInvalidInput exception when required z is missing (#1478)

Co-authored-by: Meghan Jones <[email protected]>
Co-authored-by: Dongdong Tian <[email protected]>

Author: 3 people

pygmt/clib/session.py

@@ -1362,7 +1362,14 @@ def virtualfile_from_grid(self, grid):
 
     @fmt_docstring
     def virtualfile_from_data(
-        self, check_kind=None, data=None, x=None, y=None, z=None, extra_arrays=None
+        self,
+        check_kind=None,
+        data=None,
+        x=None,
+        y=None,
+        z=None,
+        extra_arrays=None,
+        required_z=False,
     ):
         """
         Store any data inside a virtual file.
@@ -1415,7 +1422,7 @@ def virtualfile_from_data(
         ...
         <vector memory>: N = 3 <7/9> <4/6> <1/3>
         """
-        kind = data_kind(data, x, y, z)
+        kind = data_kind(data, x, y, z, required_z=required_z)
 
         if check_kind == "raster" and kind not in ("file", "grid"):
             raise GMTInvalidInput(f"Unrecognized data type for grid: {type(data)}")

pygmt/helpers/utils.py

@@ -15,7 +15,7 @@
 from pygmt.exceptions import GMTInvalidInput
 
 
-def data_kind(data, x=None, y=None, z=None):
+def data_kind(data, x=None, y=None, z=None, required_z=False):
     """
     Check what kind of data is provided to a module.
 
@@ -69,7 +69,9 @@ def data_kind(data, x=None, y=None, z=None):
     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 provided both x and y.")
+        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"
@@ -78,6 +80,8 @@ def data_kind(data, x=None, y=None, z=None):
     elif hasattr(data, "__geo_interface__"):
         kind = "geojson"
     elif data is not None:
+        if required_z and data.shape[1] < 3:
+            raise GMTInvalidInput("data must provide x, y, and z columns.")
         kind = "matrix"
     else:
         kind = "vectors"

pygmt/src/blockm.py

@@ -43,7 +43,7 @@ def _blockm(block_method, table, outfile, x, y, z, **kwargs):
         with Session() as lib:
             # Choose how data will be passed into the module
             table_context = lib.virtualfile_from_data(
-                check_kind="vector", data=table, x=x, y=y, z=z
+                check_kind="vector", data=table, x=x, y=y, z=z, required_z=True
             )
             # Run blockm* on data table
             with table_context as infile:

pygmt/src/contour.py

@@ -3,7 +3,6 @@
 """
 
 from pygmt.clib import Session
-from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
     build_arg_string,
     data_kind,
@@ -127,9 +126,7 @@ def contour(self, x=None, y=None, z=None, data=None, **kwargs):
     """
     kwargs = self._preprocess(**kwargs)  # pylint: disable=protected-access
 
-    kind = data_kind(data, x, y, z)
-    if kind == "vectors" and z is None:
-        raise GMTInvalidInput("Must provided both x, y, and z.")
+    kind = data_kind(data, x, y, z, required_z=True)
 
     with Session() as lib:
         # Choose how data will be passed in to the module

pygmt/src/plot3d.py

@@ -235,7 +235,13 @@ def plot3d(
     with Session() as lib:
         # Choose how data will be passed in to the module
         file_context = lib.virtualfile_from_data(
-            check_kind="vector", data=data, x=x, y=y, z=z, extra_arrays=extra_arrays
+            check_kind="vector",
+            data=data,
+            x=x,
+            y=y,
+            z=z,
+            extra_arrays=extra_arrays,
+            required_z=True,
         )
 
         with file_context as fname:

pygmt/src/surface.py

@@ -94,9 +94,7 @@ def surface(x=None, y=None, z=None, data=None, **kwargs):
         - None if ``outgrid`` is set (grid output will be stored in file set by
           ``outgrid``)
     """
-    kind = data_kind(data, x, y, z)
-    if kind == "vectors" and z is None:
-        raise GMTInvalidInput("Must provide z with x and y.")
+    kind = data_kind(data, x, y, z, required_z=True)
 
     with GMTTempFile(suffix=".nc") as tmpfile:
         with Session() as lib:

pygmt/src/wiggle.py

@@ -107,7 +107,7 @@ def wiggle(self, x=None, y=None, z=None, data=None, **kwargs):
     with Session() as lib:
         # Choose how data will be passed in to the module
         file_context = lib.virtualfile_from_data(
-            check_kind="vector", data=data, x=x, y=y, z=z
+            check_kind="vector", data=data, x=x, y=y, z=z, required_z=True
         )
 
         with file_context as fname:

pygmt/tests/test_clib.py

@@ -1,9 +1,11 @@
 # pylint: disable=protected-access
+# pylint: disable=redefined-outer-name
 """
 Test the wrappers for the C API.
 """
 import os
 from contextlib import contextmanager
+from itertools import product
 
 import numpy as np
 import numpy.testing as npt
@@ -23,11 +25,20 @@
 from pygmt.helpers import GMTTempFile
 
 TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
+POINTS_DATA = os.path.join(TEST_DATA_DIR, "points.txt")
 
 with clib.Session() as _lib:
     gmt_version = Version(_lib.info["version"])
 
 
+@pytest.fixture(scope="module")
+def data():
+    """
+    Load the point data from the test file.
+    """
+    return np.loadtxt(POINTS_DATA)
+
+
 @contextmanager
 def mock(session, func, returns=None, mock_func=None):
     """
@@ -410,6 +421,58 @@ def test_virtual_file_bad_direction():
                 print("This should have failed")
 
 
+def test_virtualfile_from_data_required_z_matrix():
+    """
+    Test that function fails when third z column in a matrix is needed but not
+    provided.
+    """
+    data = np.ones((5, 2))
+    with clib.Session() as lib:
+        with pytest.raises(GMTInvalidInput):
+            with lib.virtualfile_from_data(data=data, required_z=True):
+                pass
+
+
+def test_virtualfile_from_data_fail_non_valid_data(data):
+    """
+    Should raise an exception if too few or too much data is given.
+    """
+    # Test all combinations where at least one data variable
+    # is not given in the x, y case:
+    for variable in product([None, data[:, 0]], repeat=2):
+        # Filter one valid configuration:
+        if not any(item is None for item in variable):
+            continue
+        with clib.Session() as lib:
+            with pytest.raises(GMTInvalidInput):
+                lib.virtualfile_from_data(
+                    x=variable[0],
+                    y=variable[1],
+                )
+
+    # Test all combinations where at least one data variable
+    # is not given in the x, y, z case:
+    for variable in product([None, data[:, 0]], repeat=3):
+        # Filter one valid configuration:
+        if not any(item is None for item in variable):
+            continue
+        with clib.Session() as lib:
+            with pytest.raises(GMTInvalidInput):
+                lib.virtualfile_from_data(
+                    x=variable[0], y=variable[1], z=variable[2], required_z=True
+                )
+
+    # Should also fail if given too much data
+    with clib.Session() as lib:
+        with pytest.raises(GMTInvalidInput):
+            lib.virtualfile_from_data(
+                x=data[:, 0],
+                y=data[:, 1],
+                z=data[:, 2],
+                data=data,
+            )
+
+
 def test_virtualfile_from_vectors():
     """
     Test the automation for transforming vectors to virtual file dataset.

pygmt/tests/test_contour.py

@@ -3,12 +3,10 @@
 Tests contour.
 """
 import os
-from itertools import product
 
 import numpy as np
 import pytest
 from pygmt import Figure
-from pygmt.exceptions import GMTInvalidInput
 
 TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
 POINTS_DATA = os.path.join(TEST_DATA_DIR, "points.txt")
@@ -30,46 +28,6 @@ def region():
     return [10, 70, -5, 10]
 
 
-def test_contour_fail_no_data(data):
-    """
-    Should raise an exception if no data is given.
-    """
-    # Contour should raise an exception if no or not sufficient data
-    # is given
-    fig = Figure()
-    # Test all combinations where at least one data variable
-    # is not given:
-    for variable in product([None, data[:, 0]], repeat=3):
-        # Filter one valid configuration:
-        if not any(item is None for item in variable):
-            continue
-        with pytest.raises(GMTInvalidInput):
-            fig.contour(
-                x=variable[0],
-                y=variable[1],
-                z=variable[2],
-                region=region,
-                projection="X4i",
-                color="red",
-                frame="afg",
-                pen="",
-            )
-    # Should also fail if given too much data
-    with pytest.raises(GMTInvalidInput):
-        fig.contour(
-            x=data[:, 0],
-            y=data[:, 1],
-            z=data[:, 2],
-            data=data,
-            region=region,
-            projection="X10c",
-            style="c0.2c",
-            color="red",
-            frame="afg",
-            pen=True,
-        )
-
-
 @pytest.mark.mpl_image_compare
 def test_contour_vec(region):
     """

pygmt/tests/test_plot3d.py

@@ -68,48 +68,6 @@ def test_plot3d_red_circles_zsize(data, region):
     return fig
 
 
-def test_plot3d_fail_no_data(data, region):
-    """
-    Plot should raise an exception if no data is given.
-    """
-    fig = Figure()
-    with pytest.raises(GMTInvalidInput):
-        fig.plot3d(
-            region=region, projection="X10c", style="c0.2c", color="red", frame="afg"
-        )
-    with pytest.raises(GMTInvalidInput):
-        fig.plot3d(
-            x=data[:, 0],
-            region=region,
-            projection="X10c",
-            style="c0.2c",
-            color="red",
-            frame="afg",
-        )
-    with pytest.raises(GMTInvalidInput):
-        fig.plot3d(
-            y=data[:, 0],
-            region=region,
-            projection="X10c",
-            style="c0.2c",
-            color="red",
-            frame="afg",
-        )
-    # Should also fail if given too much data
-    with pytest.raises(GMTInvalidInput):
-        fig.plot3d(
-            x=data[:, 0],
-            y=data[:, 1],
-            z=data[:, 2],
-            data=data,
-            region=region,
-            projection="X10c",
-            style="c0.2c",
-            color="red",
-            frame="afg",
-        )
-
-
 def test_plot3d_fail_1d_array_with_data(data, region):
     """
     Should raise an exception if array color, size, intensity and transparency