From 39bfa71980470d2225143498317f74c5e8071e73 Mon Sep 17 00:00:00 2001
From: Wei Ji <23487320+weiji14@users.noreply.github.com>
Date: Sun, 19 Sep 2021 12:09:59 +1200
Subject: [PATCH 1/3] Let required_z check work on xarray.Dataset inputs

Modify if-statement check in `data_kind` helper function
to use `len(data.data_vars)` to check the shape of
xarray.Dataset inputs. Added two regression tests in
test_clib and test_blockm to ensure this works on both the
low-level clib and high-level module APIs.
---
 pygmt/helpers/utils.py     |  5 ++++-
 pygmt/tests/test_blockm.py |  7 +++++--
 pygmt/tests/test_clib.py   | 31 ++++++++++++++++++++++++++++++-
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/pygmt/helpers/utils.py b/pygmt/helpers/utils.py
index 19b6e2ec2f7..63eab2cb996 100644
--- a/pygmt/helpers/utils.py
+++ b/pygmt/helpers/utils.py
@@ -80,7 +80,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 data.shape[1] < 3:
+        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:
diff --git a/pygmt/tests/test_blockm.py b/pygmt/tests/test_blockm.py
index 39a0e315f7f..d71b372b6b2 100644
--- a/pygmt/tests/test_blockm.py
+++ b/pygmt/tests/test_blockm.py
@@ -3,9 +3,11 @@
 """
 import os
 
+import numpy as np
 import numpy.testing as npt
 import pandas as pd
 import pytest
+import xarray as xr
 from pygmt import blockmean, blockmode
 from pygmt.datasets import load_sample_bathymetry
 from pygmt.exceptions import GMTInvalidInput
@@ -31,12 +33,13 @@ def test_blockmean_input_dataframe(dataframe):
     npt.assert_allclose(output.iloc[0], [245.888877, 29.978707, -384.0])
 
 
-def test_blockmean_input_table_matrix(dataframe):
+@pytest.mark.parametrize("array_func", [np.array, xr.Dataset])
+def test_blockmean_input_table_matrix(array_func, dataframe):
     """
     Run blockmean using table input that is not a pandas.DataFrame but still a
     matrix.
     """
-    table = dataframe.values
+    table = array_func(dataframe)
     output = blockmean(table=table, spacing="5m", region=[245, 255, 20, 30])
     assert isinstance(output, pd.DataFrame)
     assert output.shape == (5849, 3)
diff --git a/pygmt/tests/test_clib.py b/pygmt/tests/test_clib.py
index 2be988ec487..22fd4d6d7b8 100644
--- a/pygmt/tests/test_clib.py
+++ b/pygmt/tests/test_clib.py
@@ -419,7 +419,36 @@ def test_virtual_file_bad_direction():
                 print("This should have failed")
 
 
-def test_virtualfile_from_data_required_z_matrix():
+@pytest.mark.parametrize(
+    "array_func,kind",
+    [(np.array, "matrix"), (pd.DataFrame, "vector"), (xr.Dataset, "vector")],
+)
+def test_virtualfile_from_data_required_z_matrix(array_func, kind):
+    """
+    Test that function works when third z column in a matrix is needed and
+    provided.
+    """
+    shape = (5, 3)
+    dataframe = pd.DataFrame(
+        data=np.arange(shape[0] * shape[1]).reshape(shape), columns=["x", "y", "z"]
+    )
+    data = array_func(dataframe)
+    with clib.Session() as lib:
+        with lib.virtualfile_from_data(data=data, required_z=True) as vfile:
+            with GMTTempFile() as outfile:
+                lib.call_module("info", f"{vfile} ->{outfile.name}")
+                output = outfile.read(keep_tabs=True)
+        bounds = "\t".join(
+            [
+                f"<{i.min():.0f}/{i.max():.0f}>"
+                for i in (dataframe.x, dataframe.y, dataframe.z)
+            ]
+        )
+        expected = f"<{kind} memory>: N = {shape[0]}\t{bounds}\n"
+        assert output == expected
+
+
+def test_virtualfile_from_data_required_z_matrix_missing():
     """
     Test that function fails when third z column in a matrix is needed but not
     provided.

From 4efa5135463962064365137edaad385721f82f04 Mon Sep 17 00:00:00 2001
From: Wei Ji <23487320+weiji14@users.noreply.github.com>
Date: Sun, 19 Sep 2021 12:24:53 +1200
Subject: [PATCH 2/3] Mention required_z parameter in docstring

---
 pygmt/clib/session.py  | 2 ++
 pygmt/helpers/utils.py | 5 +++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py
index 8d7ee8d71ed..569c01ede25 100644
--- a/pygmt/clib/session.py
+++ b/pygmt/clib/session.py
@@ -1375,6 +1375,8 @@ def virtualfile_from_data(
         extra_arrays : list of 1d arrays
             Optional. A list of numpy arrays in addition to x, y and z. All
             of these arrays must be of the same size as the x/y/z arrays.
+        required_z : bool
+            State whether the 'z' column is required. [Default is False].
 
         Returns
         -------
diff --git a/pygmt/helpers/utils.py b/pygmt/helpers/utils.py
index 63eab2cb996..15d06be992f 100644
--- a/pygmt/helpers/utils.py
+++ b/pygmt/helpers/utils.py
@@ -39,8 +39,9 @@ def data_kind(data, x=None, y=None, z=None, required_z=False):
     x/y : 1d arrays or None
         x and y columns as numpy arrays.
     z : 1d array or None
-        z column as numpy array. To be used optionally when x and y
-        are given.
+        z column as numpy array. To be used optionally when x and y are given.
+    required_z : bool
+        State whether the 'z' column is required. [Default is False].
 
     Returns
     -------

From 3f1e6723a8d0ebe9eec686c17a03e36a97267a64 Mon Sep 17 00:00:00 2001
From: Wei Ji <23487320+weiji14@users.noreply.github.com>
Date: Sun, 19 Sep 2021 23:31:48 +1200
Subject: [PATCH 3/3] Apply suggestions from code review

Co-authored-by: Dongdong Tian <seisman.info@gmail.com>
---
 pygmt/clib/session.py  | 2 +-
 pygmt/helpers/utils.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py
index 569c01ede25..77925597e3f 100644
--- a/pygmt/clib/session.py
+++ b/pygmt/clib/session.py
@@ -1376,7 +1376,7 @@ def virtualfile_from_data(
             Optional. A list of numpy arrays in addition to x, y and z. All
             of these arrays must be of the same size as the x/y/z arrays.
         required_z : bool
-            State whether the 'z' column is required. [Default is False].
+            State whether the 'z' column is required.
 
         Returns
         -------
diff --git a/pygmt/helpers/utils.py b/pygmt/helpers/utils.py
index 15d06be992f..2f3ce181164 100644
--- a/pygmt/helpers/utils.py
+++ b/pygmt/helpers/utils.py
@@ -41,7 +41,7 @@ def data_kind(data, x=None, y=None, z=None, required_z=False):
     z : 1d array or None
         z column as numpy array. To be used optionally when x and y are given.
     required_z : bool
-        State whether the 'z' column is required. [Default is False].
+        State whether the 'z' column is required.
 
     Returns
     -------