Raise GMTInvalidInput exception when required z is missing

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.
@@ -1414,7 +1421,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

@@ -14,7 +14,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.
 
@@ -65,6 +65,8 @@ def data_kind(data, x=None, y=None, z=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.")
+    if data is None and required_z and (x is None or y is None or z is None):
+        raise GMTInvalidInput("Must provided both x, y, and z.")
 
     if isinstance(data, str):
         kind = "file"

pygmt/tests/test_clib.py

@@ -410,6 +410,19 @@ def test_virtual_file_bad_direction():
                 print("This should have failed")
 
 
+def test_virtualfile_from_data_required_z():
+    """
+    Test that function fails when needs and do not provide z columns.
+    """
+    x = np.arange(5)
+    y = np.arange(5)
+    z = None
+    with clib.Session() as lib:
+        with pytest.raises(GMTInvalidInput):
+            with lib.virtualfile_from_data(x=x, y=y, z=z, required_z=True):
+                print("This should have failed")
+
+
 def test_virtualfile_from_vectors():
     """
     Test the automation for transforming vectors to virtual file dataset.