add formatting to return array or DataFrame

Author: willschlitzer

pygmt/src/grdvolume.py

@@ -1,7 +1,10 @@
 """
 grdvolume - Calculate grid volume and area constrained by a contour.
 """
+import numpy as np
+import pandas as pd
 from pygmt.clib import Session
+from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
     GMTTempFile,
     build_arg_string,
@@ -21,7 +24,7 @@
     V="verbose",
 )
 @kwargs_to_strings(C="sequence", R="sequence")
-def grdvolume(grid, **kwargs):
+def grdvolume(grid, data_format="a", **kwargs):
     r"""
     {aliases}
 
@@ -47,4 +50,20 @@ def grdvolume(grid, **kwargs):
                 )
                 lib.call_module("grdvolume", arg_str)
         result = outfile.read()
+    if data_format == "s":
+        return result
+    data_list = []
+    for string_entry in result.strip().split("\n"):
+        float_entry = []
+        string_list = string_entry.strip().split()
+        for i in string_list:
+            float_entry.append(float(i))
+        data_list.append(float_entry)
+    data_array = np.array(data_list)
+    if data_format == "a":
+        result = data_array
+    elif data_format == "d":
+        result = pd.DataFrame(data_array)
+    else:
+        raise GMTInvalidInput("""Must specify format as either a, d, or s.""")
     return result