Make all functions/methods have consistent behavior for table output

Author: seisman

pygmt/src/blockm.py

@@ -6,17 +6,17 @@
 import pandas as pd
 from pygmt.clib import Session
 from pygmt.helpers import (
-    GMTTempFile,
     build_arg_string,
     fmt_docstring,
     kwargs_to_strings,
     use_alias,
+    validate_output_table_type,
 )
 
 __doctest_skip__ = ["blockmean", "blockmedian", "blockmode"]
 
 
-def _blockm(block_method, data, x, y, z, outfile, **kwargs):
+def _blockm(block_method, data, x, y, z, output_type, outfile, **kwargs):
     r"""
     Block average (x, y, z) data tables by mean, median, or mode estimation.
 
@@ -42,30 +42,28 @@ def _blockm(block_method, data, x, y, z, outfile, **kwargs):
         - None if ``outfile`` is set (filtered output will be stored in file
           set by ``outfile``)
     """
-    with GMTTempFile(suffix=".csv") as tmpfile:
-        with Session() as lib:
-            with lib.virtualfile_in(
+    output_type = validate_output_table_type(output_type, outfile=outfile)
+
+    with Session() as lib:
+        with (
+            lib.virtualfile_in(
                 check_kind="vector", data=data, x=x, y=y, z=z, required_z=True
-            ) as vintbl:
-                # Run blockm* on data table
-                if outfile is None:
-                    outfile = tmpfile.name
-                lib.call_module(
-                    module=block_method,
-                    args=build_arg_string(kwargs, infile=vintbl, outfile=outfile),
-                )
-
-        # Read temporary csv output to a pandas table
-        if outfile == tmpfile.name:  # if user did not set outfile, return pd.DataFrame
-            try:
-                column_names = data.columns.to_list()
-                result = pd.read_csv(tmpfile.name, sep="\t", names=column_names)
-            except AttributeError:  # 'str' object has no attribute 'columns'
-                result = pd.read_csv(tmpfile.name, sep="\t", header=None, comment=">")
-        elif outfile != tmpfile.name:  # return None if outfile set, output in outfile
-            result = None
-
-    return result
+            ) as vintbl,
+            lib.virtualfile_out(kind="dataset", fname=outfile) as vouttbl,
+        ):
+            lib.call_module(
+                module=block_method,
+                args=build_arg_string(kwargs, infile=vintbl, outfile=vouttbl),
+            )
+        column_names = None
+        if isinstance(data, pd.DataFrame):
+            column_names = data.columns.to_list()
+
+        return lib.return_table(
+            output_type=output_type,
+            vfile=vouttbl,
+            column_names=column_names,
+        )
 
 
 @fmt_docstring
@@ -86,7 +84,9 @@ def _blockm(block_method, data, x, y, z, outfile, **kwargs):
     w="wrap",
 )
 @kwargs_to_strings(I="sequence", R="sequence", i="sequence_comma", o="sequence_comma")
-def blockmean(data=None, x=None, y=None, z=None, outfile=None, **kwargs):
+def blockmean(
+    data=None, x=None, y=None, z=None, output_type="pandas", outfile=None, **kwargs
+):
     r"""
     Block average (x, y, z) data tables by mean estimation.
 
@@ -159,7 +159,14 @@ def blockmean(data=None, x=None, y=None, z=None, outfile=None, **kwargs):
     >>> data_bmean = pygmt.blockmean(data=data, region=[245, 255, 20, 30], spacing="5m")
     """
     return _blockm(
-        block_method="blockmean", data=data, x=x, y=y, z=z, outfile=outfile, **kwargs
+        block_method="blockmean",
+        data=data,
+        x=x,
+        y=y,
+        z=z,
+        output_type=output_type,
+        outfile=outfile,
+        **kwargs,
     )
 
 
@@ -180,7 +187,9 @@ def blockmean(data=None, x=None, y=None, z=None, outfile=None, **kwargs):
     w="wrap",
 )
 @kwargs_to_strings(I="sequence", R="sequence", i="sequence_comma", o="sequence_comma")
-def blockmedian(data=None, x=None, y=None, z=None, outfile=None, **kwargs):
+def blockmedian(
+    data=None, x=None, y=None, z=None, output_type="pandas", outfile=None, **kwargs
+):
     r"""
     Block average (x, y, z) data tables by median estimation.
 
@@ -246,7 +255,14 @@ def blockmedian(data=None, x=None, y=None, z=None, outfile=None, **kwargs):
     ... )
     """
     return _blockm(
-        block_method="blockmedian", data=data, x=x, y=y, z=z, outfile=outfile, **kwargs
+        block_method="blockmedian",
+        data=data,
+        x=x,
+        y=y,
+        z=z,
+        output_type=output_type,
+        outfile=outfile,
+        **kwargs,
     )
 
 
@@ -267,7 +283,9 @@ def blockmedian(data=None, x=None, y=None, z=None, outfile=None, **kwargs):
     w="wrap",
 )
 @kwargs_to_strings(I="sequence", R="sequence", i="sequence_comma", o="sequence_comma")
-def blockmode(data=None, x=None, y=None, z=None, outfile=None, **kwargs):
+def blockmode(
+    data=None, x=None, y=None, z=None, output_type="pandas", outfile=None, **kwargs
+):
     r"""
     Block average (x, y, z) data tables by mode estimation.
 
@@ -331,5 +349,12 @@ def blockmode(data=None, x=None, y=None, z=None, outfile=None, **kwargs):
     >>> data_bmode = pygmt.blockmode(data=data, region=[245, 255, 20, 30], spacing="5m")
     """
     return _blockm(
-        block_method="blockmode", data=data, x=x, y=y, z=z, outfile=outfile, **kwargs
+        block_method="blockmode",
+        data=data,
+        x=x,
+        y=y,
+        z=z,
+        output_type=output_type,
+        outfile=outfile,
+        **kwargs,
     )

pygmt/src/filter1d.py

@@ -2,11 +2,9 @@
 filter1d - Time domain filtering of 1-D data tables
 """
 
-import pandas as pd
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
-    GMTTempFile,
     build_arg_string,
     fmt_docstring,
     use_alias,
@@ -117,22 +115,13 @@ def filter1d(data, output_type="pandas", outfile=None, **kwargs):
 
     output_type = validate_output_table_type(output_type, outfile=outfile)
 
-    with GMTTempFile() as tmpfile:
-        with Session() as lib:
-            with lib.virtualfile_in(check_kind="vector", data=data) as vintbl:
-                if outfile is None:
-                    outfile = tmpfile.name
-                lib.call_module(
-                    module="filter1d",
-                    args=build_arg_string(kwargs, infile=vintbl, outfile=outfile),
-                )
-
-        # Read temporary csv output to a pandas table
-        if outfile == tmpfile.name:  # if user did not set outfile, return pd.DataFrame
-            result = pd.read_csv(tmpfile.name, sep="\t", header=None, comment=">")
-        elif outfile != tmpfile.name:  # return None if outfile set, output in outfile
-            result = None
-
-        if output_type == "numpy":
-            result = result.to_numpy()
-    return result
+    with Session() as lib:
+        with (
+            lib.virtualfile_in(check_kind="vector", data=data) as vintbl,
+            lib.virtualfile_out(kind="dataset", fname=outfile) as vouttbl,
+        ):
+            lib.call_module(
+                module="filter1d",
+                args=build_arg_string(kwargs, infile=vintbl, outfile=vouttbl),
+            )
+        return lib.return_table(output_type=output_type, vfile=vouttbl)

pygmt/src/grd2xyz.py

@@ -2,12 +2,10 @@
 grd2xyz - Convert grid to data table
 """
 
-import pandas as pd
 import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
-    GMTTempFile,
     build_arg_string,
     fmt_docstring,
     kwargs_to_strings,
@@ -150,30 +148,23 @@ def grd2xyz(grid, output_type="pandas", outfile=None, **kwargs):
         )
 
     # Set the default column names for the pandas dataframe header
-    dataframe_header = ["x", "y", "z"]
+    column_names = ["x", "y", "z"]
     # Let output pandas column names match input DataArray dimension names
-    if isinstance(grid, xr.DataArray) and output_type == "pandas":
+    if isinstance(grid, xr.DataArray):
         # Reverse the dims because it is rows, columns ordered.
-        dataframe_header = [grid.dims[1], grid.dims[0], grid.name]
-
-    with GMTTempFile() as tmpfile:
-        with Session() as lib:
-            with lib.virtualfile_in(check_kind="raster", data=grid) as vingrd:
-                if outfile is None:
-                    outfile = tmpfile.name
-                lib.call_module(
-                    module="grd2xyz",
-                    args=build_arg_string(kwargs, infile=vingrd, outfile=outfile),
-                )
-
-        # Read temporary csv output to a pandas table
-        if outfile == tmpfile.name:  # if user did not set outfile, return pd.DataFrame
-            result = pd.read_csv(
-                tmpfile.name, sep="\t", names=dataframe_header, comment=">"
+        column_names = [grid.dims[1], grid.dims[0], grid.name]
+
+    with Session() as lib:
+        with (
+            lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,
+            lib.virtualfile_out(kind="dataset", fname=outfile) as vouttbl,
+        ):
+            lib.call_module(
+                module="grd2xyz",
+                args=build_arg_string(kwargs, infile=vingrd, outfile=vouttbl),
             )
-        elif outfile != tmpfile.name:  # return None if outfile set, output in outfile
-            result = None
-
-        if output_type == "numpy":
-            result = result.to_numpy()
-    return result
+        return lib.return_table(
+            output_type=output_type,
+            vfile=vouttbl,
+            column_names=column_names,
+        )

pygmt/src/grdhisteq.py

@@ -3,7 +3,6 @@
 """
 
 import numpy as np
-import pandas as pd
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -231,33 +230,28 @@ def compute_bins(grid, output_type="pandas", **kwargs):
         if kwargs.get("h") is not None and output_type != "file":
             raise GMTInvalidInput("'header' is only allowed with output_type='file'.")
 
-        with GMTTempFile(suffix=".txt") as tmpfile:
-            with Session() as lib:
-                with lib.virtualfile_in(check_kind="raster", data=grid) as vingrd:
-                    if outfile is None:
-                        kwargs["D"] = outfile = tmpfile.name  # output to tmpfile
-                    lib.call_module(
-                        module="grdhisteq", args=build_arg_string(kwargs, infile=vingrd)
-                    )
+        with Session() as lib:
+            with (
+                lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,
+                lib.virtualfile_out(kind="dataset", fname=outfile) as vouttbl,
+            ):
+                kwargs["D"] = vouttbl  # -D for output file name
+                lib.call_module(
+                    module="grdhisteq", args=build_arg_string(kwargs, infile=vingrd)
+                )
 
-            if outfile == tmpfile.name:
-                # if user did not set outfile, return pd.DataFrame
-                result = pd.read_csv(
-                    filepath_or_buffer=outfile,
-                    sep="\t",
-                    header=None,
-                    names=["start", "stop", "bin_id"],
-                    dtype={
+            result = lib.return_table(
+                output_type=output_type,
+                vfile=vouttbl,
+                column_names=["start", "stop", "bin_id"],
+            )
+            if output_type == "pandas":
+                result = result.astype(
+                    {
                         "start": np.float32,
                         "stop": np.float32,
                         "bin_id": np.uint32,
-                    },
+                    }
                 )
-            elif outfile != tmpfile.name:
-                # return None if outfile set, output in outfile
-                return None
-
-            if output_type == "numpy":
-                return result.to_numpy()
-
-            return result.set_index("bin_id")
+                return result.set_index("bin_id")
+            return result

pygmt/src/grdtrack.py

@@ -6,11 +6,11 @@
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
-    GMTTempFile,
     build_arg_string,
     fmt_docstring,
     kwargs_to_strings,
     use_alias,
+    validate_output_table_type,
 )
 
 __doctest_skip__ = ["grdtrack"]
@@ -44,7 +44,9 @@
     w="wrap",
 )
 @kwargs_to_strings(R="sequence", S="sequence", i="sequence_comma", o="sequence_comma")
-def grdtrack(grid, points=None, newcolname=None, outfile=None, **kwargs):
+def grdtrack(
+    grid, points=None, output_type="pandas", outfile=None, newcolname=None, **kwargs
+):
     r"""
     Sample grids at specified (x,y) locations.
 
@@ -291,30 +293,27 @@ def grdtrack(grid, points=None, newcolname=None, outfile=None, **kwargs):
     if hasattr(points, "columns") and newcolname is None:
         raise GMTInvalidInput("Please pass in a str to 'newcolname'")
 
-    with GMTTempFile(suffix=".csv") as tmpfile:
-        with Session() as lib:
-            with (
-                lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,
-                lib.virtualfile_in(
-                    check_kind="vector", data=points, required_data=False
-                ) as vintbl,
-            ):
-                kwargs["G"] = vingrd
-                if outfile is None:  # Output to tmpfile if outfile is not set
-                    outfile = tmpfile.name
-                lib.call_module(
-                    module="grdtrack",
-                    args=build_arg_string(kwargs, infile=vintbl, outfile=outfile),
-                )
+    output_type = validate_output_table_type(output_type, outfile=outfile)
 
-        # Read temporary csv output to a pandas table
-        if outfile == tmpfile.name:  # if user did not set outfile, return pd.DataFrame
-            try:
-                column_names = [*points.columns.to_list(), newcolname]
-                result = pd.read_csv(tmpfile.name, sep="\t", names=column_names)
-            except AttributeError:  # 'str' object has no attribute 'columns'
-                result = pd.read_csv(tmpfile.name, sep="\t", header=None, comment=">")
-        elif outfile != tmpfile.name:  # return None if outfile set, output in outfile
-            result = None
+    column_names = None
+    if isinstance(points, pd.DataFrame):
+        column_names = [*points.columns.to_list(), newcolname]
 
-    return result
+    with Session() as lib:
+        with (
+            lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,
+            lib.virtualfile_in(
+                check_kind="vector", data=points, required_data=False
+            ) as vintbl,
+            lib.virtualfile_out(kind="dataset", fname=outfile) as vouttbl,
+        ):
+            kwargs["G"] = vingrd
+            lib.call_module(
+                module="grdtrack",
+                args=build_arg_string(kwargs, infile=vintbl, outfile=vouttbl),
+            )
+        return lib.return_table(
+            output_type=output_type,
+            vfile=vouttbl,
+            column_names=column_names,
+        )