Add pygmt.load_dataarray function and refactor modules that return None or xarray.dataarray

pygmt/io.py

@@ -0,0 +1,36 @@
+"""
+PyGMT input/output (I/O) utilities.
+"""
+import xarray as xr
+
+
+def process_output_grid(grid_name, tmpfile_name):
+    """
+    Processes the output from the GMT API to return an xarray.DataArray if
+    ``grid_name`` matches ``tmpfile_name`` and return None if it does not.
+
+    Parameters
+    ----------
+    grid_name : str
+        The name of the output netCDF file with extension .nc to store the grid
+        in.
+    tmpfile_name : str
+        The name attribute from a GMTTempFile instance.
+
+    Returns
+    -------
+    ret: xarray.DataArray or None
+        Return type depends on whether the ``outgrid`` parameter is set:
+
+        - :class:`xarray.DataArray` if ``outgrid`` is not set
+        - None if ``outgrid`` is set (grid output will be stored in file set by
+          ``outgrid``)
+    """
+    if grid_name == tmpfile_name:  # Implies user did not set outgrid, return DataArray
+        with xr.open_dataarray(grid_name) as dataarray:
+            result = dataarray.load()
+            _ = result.gmt  # load GMTDataArray accessor information
+    else:
+        result = None  # Implies user set an outgrid, return None
+
+    return result

pygmt/src/grdclip.py

@@ -2,7 +2,6 @@
 grdclip - Change the range and extremes of grid values.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -11,6 +10,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -88,11 +88,4 @@ def grdclip(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdclip", arg_str)
 
-        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
-            with xr.open_dataarray(outgrid) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return process_output_grid(outgrid, tmpfile.name)

pygmt/src/grdcut.py

@@ -2,7 +2,6 @@
 grdcut - Extract subregion from a grid.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -11,6 +10,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -98,11 +98,4 @@ def grdcut(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdcut", arg_str)
 
-        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
-            with xr.open_dataarray(outgrid) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return process_output_grid(outgrid, tmpfile.name)

pygmt/src/grdfill.py

@@ -2,7 +2,6 @@
 grdfill - Fill blank areas from a grid.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -12,6 +11,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -75,11 +75,4 @@ def grdfill(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdfill", arg_str)
 
-        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
-            with xr.open_dataarray(outgrid) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return process_output_grid(outgrid, tmpfile.name)

pygmt/src/grdfilter.py

@@ -2,7 +2,6 @@
 grdfilter - Filter a grid in the space (or time) domain.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -11,6 +10,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -151,11 +151,4 @@ def grdfilter(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdfilter", arg_str)
 
-        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
-            with xr.open_dataarray(outgrid) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return process_output_grid(outgrid, tmpfile.name)

pygmt/src/grdgradient.py

@@ -2,7 +2,6 @@
 grdgradient - Compute directional gradients from a grid.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -13,6 +12,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -117,11 +117,4 @@ def grdgradient(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdgradient", arg_str)
 
-        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
-            with xr.open_dataarray(outgrid) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return process_output_grid(outgrid, tmpfile.name)

pygmt/src/grdlandmask.py

@@ -2,7 +2,6 @@
 grdlandmask - Create a "wet-dry" mask grid from shoreline data base
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -12,6 +11,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -68,11 +68,4 @@ def grdlandmask(**kwargs):
             arg_str = build_arg_string(kwargs)
             lib.call_module("grdlandmask", arg_str)
 
-        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
-            with xr.open_dataarray(outgrid) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return process_output_grid(outgrid, tmpfile.name)

pygmt/src/grdsample.py

@@ -2,7 +2,6 @@
 grdsample - Resample a grid onto a new lattice
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -11,6 +10,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -85,11 +85,4 @@ def grdsample(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdsample", arg_str)
 
-        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
-            with xr.open_dataarray(outgrid) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return process_output_grid(outgrid, tmpfile.name)

pygmt/src/surface.py

@@ -2,7 +2,6 @@
 surface - Grids table data using adjustable tension continuous curvature
 splines.
 """
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -14,6 +13,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -99,11 +99,4 @@ def surface(x=None, y=None, z=None, data=None, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module(module="surface", args=arg_str)
 
-        if outfile == tmpfile.name:  # if user did not set outfile, return DataArray
-            with xr.open_dataarray(outfile) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        elif outfile != tmpfile.name:  # if user sets an outfile, return None
-            result = None
-
-    return result
+        return process_output_grid(outfile, tmpfile.name)

pygmt/src/xyz2grd.py

@@ -1,7 +1,6 @@
 """
 xyz2grd - Convert data table to a grid.
 """
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -10,6 +9,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import process_output_grid
 
 
 @fmt_docstring
@@ -64,11 +64,4 @@ def xyz2grd(table, **kwargs):
                 arg_str = " ".join([infile, arg_str])
                 lib.call_module("xyz2grd", arg_str)
 
-        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
-            with xr.open_dataarray(outgrid) as dataarray:
-                result = dataarray.load()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return process_output_grid(outgrid, tmpfile.name)