Add pygmt.read to read a dataset/grid/image into pandas.DataFrame/xarray.DataArray

doc/api/index.rst

@@ -172,6 +172,7 @@ Input/output
     :toctree: generated
 
     load_dataarray
+    read
 
 GMT Defaults
 ------------

pygmt/__init__.py

@@ -54,6 +54,7 @@
     makecpt,
     nearneighbor,
     project,
+    read,
     select,
     sph2grd,
     sphdistance,

pygmt/datasets/samples.py

@@ -8,8 +8,7 @@
 import pandas as pd
 import xarray as xr
 from pygmt.exceptions import GMTInvalidInput
-from pygmt.io import load_dataarray
-from pygmt.src import which
+from pygmt.src import read, which
 
 
 def _load_japan_quakes() -> pd.DataFrame:
@@ -203,8 +202,7 @@ def _load_earth_relief_holes() -> xr.DataArray:
         The Earth relief grid. Coordinates are latitude and longitude in degrees. Relief
         is in meters.
     """
-    fname = which("@earth_relief_20m_holes.grd", download="c")
-    return load_dataarray(fname, engine="netcdf4")
+    return read("@earth_relief_20m_holes.grd", kind="grid")  # type: ignore[return-value]
 
 
 class GMTSampleData(NamedTuple):

pygmt/helpers/testing.py

@@ -7,9 +7,9 @@
 import string
 from pathlib import Path
 
+import xarray as xr
 from pygmt.exceptions import GMTImageComparisonFailure
-from pygmt.io import load_dataarray
-from pygmt.src import which
+from pygmt.src import read
 
 
 def check_figures_equal(*, extensions=("png",), tol=0.0, result_dir="result_images"):
@@ -144,17 +144,16 @@ def wrapper(*args, ext="png", request=None, **kwargs):
     return decorator
 
 
-def load_static_earth_relief():
+def load_static_earth_relief() -> xr.DataArray:
     """
-    Load the static_earth_relief file for internal testing.
+    Load the static_earth_relief.nc file for internal testing.
 
     Returns
     -------
-    data : xarray.DataArray
+    data
         A grid of Earth relief for internal tests.
     """
-    fname = which("@static_earth_relief.nc", download="c")
-    return load_dataarray(fname)
+    return read("@static_earth_relief.nc", kind="grid")  # type: ignore[return-value]
 
 
 def skip_if_no(package):

pygmt/src/__init__.py

@@ -41,6 +41,7 @@
 from pygmt.src.plot3d import plot3d
 from pygmt.src.project import project
 from pygmt.src.psconvert import psconvert
+from pygmt.src.read import read
 from pygmt.src.rose import rose
 from pygmt.src.select import select
 from pygmt.src.shift_origin import shift_origin

pygmt/src/read.py

@@ -0,0 +1,118 @@
+"""
+Read a file into an appropriate object.
+"""
+
+from collections.abc import Mapping, Sequence
+from pathlib import PurePath
+from typing import Any, Literal
+
+import pandas as pd
+import xarray as xr
+from pygmt.clib import Session
+from pygmt.helpers import build_arg_list, is_nonstr_iter
+from pygmt.src.which import which
+
+
+def read(
+    file: str | PurePath,
+    kind: Literal["dataset", "grid", "image"],
+    region: Sequence[float] | str | None = None,
+    header: int | None = None,
+    column_names: pd.Index | None = None,
+    dtype: type | Mapping[Any, type] | None = None,
+    index_col: str | int | None = None,
+) -> pd.DataFrame | xr.DataArray:
+    """
+    Read a dataset, grid, or image from a file and return the appropriate object.
+
+    The returned object is a :class:`pandas.DataFrame` for datasets, and
+    :class:`xarray.DataArray` for grids and images.
+
+    For datasets, keyword arguments ``column_names``, ``header``, ``dtype``, and
+    ``index_col`` are supported.
+
+    Parameters
+    ----------
+    file
+        The file name to read.
+    kind
+        The kind of data to read. Valid values are ``"dataset"``, ``"grid"``, and
+        ``"image"``.
+    region
+        The region of interest. Only data within this region will be read.
+    column_names
+        A list of column names.
+    header
+        Row number containing column names. ``header=None`` means not to parse the
+        column names from table header. Ignored if the row number is larger than the
+        number of headers in the table.
+    dtype
+        Data type. Can be a single type for all columns or a dictionary mapping
+        column names to types.
+    index_col
+        Column to set as index.
+
+    Returns
+    -------
+    data
+        Return type depends on the ``kind`` argument:
+
+        - ``"dataset"``: :class:`pandas.DataFrame`
+        - ``"grid"`` or ``"image"``: :class:`xarray.DataArray`
+
+
+    Examples
+    --------
+    Read a dataset into a :class:`pandas.DataFrame` object:
+
+    >>> from pygmt import read
+    >>> df = read("@hotspots.txt", kind="dataset")
+    >>> type(df)
+    <class 'pandas.core.frame.DataFrame'>
+
+    Read a grid into an :class:`xarray.DataArray` object:
+
+    >>> dataarray = read("@earth_relief_01d", kind="grid")
+    >>> type(dataarray)
+    <class 'xarray.core.dataarray.DataArray'>
+    """
+    if kind not in {"dataset", "grid", "image"}:
+        msg = f"Invalid kind {kind}: must be one of 'dataset', 'grid', or 'image'."
+        raise ValueError(msg)
+
+    if kind != "dataset" and any(
+        v is not None for v in [column_names, header, dtype, index_col]
+    ):
+        msg = (
+            "Only the 'dataset' kind supports the 'column_names', 'header', "
+            "'dtype', and 'index_col' arguments."
+        )
+        raise ValueError(msg)
+
+    kwdict = {
+        "R": "/".join(f"{v}" for v in region) if is_nonstr_iter(region) else region,  # type: ignore[union-attr]
+        "T": {"dataset": "d", "grid": "g", "image": "i"}[kind],
+    }
+
+    with Session() as lib:
+        with lib.virtualfile_out(kind=kind) as voutfile:
+            lib.call_module("read", args=[file, voutfile, *build_arg_list(kwdict)])
+
+        match kind:
+            case "dataset":
+                return lib.virtualfile_to_dataset(
+                    vfname=voutfile,
+                    column_names=column_names,
+                    header=header,
+                    dtype=dtype,
+                    index_col=index_col,
+                )
+            case "grid" | "image":
+                raster = lib.virtualfile_to_raster(vfname=voutfile, kind=kind)
+                # Add "source" encoding
+                source = which(fname=file)
+                raster.encoding["source"] = (
+                    source[0] if isinstance(source, list) else source
+                )
+                _ = raster.gmt  # Load GMTDataArray accessor information
+                return raster

pygmt/tests/test_clib_put_matrix.py

@@ -5,8 +5,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
-import xarray as xr
-from pygmt import clib
+from pygmt import clib, read
 from pygmt.clib.session import DTYPES_NUMERIC
 from pygmt.exceptions import GMTCLibError
 from pygmt.helpers import GMTTempFile
@@ -101,7 +100,7 @@ def test_put_matrix_grid(dtypes):
                 newdata = tmp_file.loadtxt(dtype=dtype)
                 npt.assert_allclose(newdata, data)
 
-            # Save the data to a netCDF grid and check that xarray can load it
+            # Save the data to a netCDF grid and check it can be read again.
             with GMTTempFile(suffix=".nc") as tmp_grid:
                 lib.write_data(
                     "GMT_IS_MATRIX",
@@ -111,12 +110,12 @@ def test_put_matrix_grid(dtypes):
                     tmp_grid.name,
                     grid,
                 )
-                with xr.open_dataarray(tmp_grid.name) as dataarray:
-                    assert dataarray.shape == shape
-                    npt.assert_allclose(dataarray.data, np.flipud(data))
-                    npt.assert_allclose(
-                        dataarray.coords["x"].actual_range, np.array(wesn[0:2])
-                    )
-                    npt.assert_allclose(
-                        dataarray.coords["y"].actual_range, np.array(wesn[2:4])
-                    )
+                dataarray = read(tmp_grid.name, kind="grid")
+                assert dataarray.shape == shape
+                npt.assert_allclose(dataarray.data, np.flipud(data))
+                npt.assert_allclose(
+                    dataarray.coords["x"].actual_range, np.array(wesn[0:2])
+                )
+                npt.assert_allclose(
+                    dataarray.coords["y"].actual_range, np.array(wesn[2:4])
+                )

pygmt/tests/test_clib_read_data.py

@@ -11,7 +11,7 @@
 from pygmt.clib import Session
 from pygmt.exceptions import GMTCLibError
 from pygmt.helpers import GMTTempFile
-from pygmt.io import load_dataarray
+from pygmt.helpers.testing import load_static_earth_relief
 from pygmt.src import which
 
 try:
@@ -27,7 +27,7 @@ def fixture_expected_xrgrid():
     """
     The expected xr.DataArray object for the static_earth_relief.nc file.
     """
-    return load_dataarray(which("@static_earth_relief.nc"))
+    return load_static_earth_relief()
 
 
 @pytest.fixture(scope="module", name="expected_xrimage")

pygmt/tests/test_datatypes_dataset.py

@@ -6,8 +6,7 @@
 
 import pandas as pd
 import pytest
-from pygmt import which
-from pygmt.clib import Session
+from pygmt import read, which
 from pygmt.helpers import GMTTempFile
 
 
@@ -44,11 +43,7 @@ def dataframe_from_gmt(fname, **kwargs):
     """
     Read tabular data as pandas.DataFrame using GMT virtual file.
     """
-    with Session() as lib:
-        with lib.virtualfile_out(kind="dataset") as vouttbl:
-            lib.call_module("read", [fname, vouttbl, "-Td"])
-            df = lib.virtualfile_to_dataset(vfname=vouttbl, **kwargs)
-            return df
+    return read(fname, kind="dataset", **kwargs)
 
 
 @pytest.mark.benchmark

pygmt/tests/test_dimfilter.py

@@ -6,7 +6,7 @@
 
 import pytest
 import xarray as xr
-from pygmt import dimfilter, load_dataarray
+from pygmt import dimfilter, read
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import GMTTempFile
 from pygmt.helpers.testing import load_static_earth_relief
@@ -56,7 +56,7 @@ def test_dimfilter_outgrid(grid, expected_grid):
         )
         assert result is None  # return value is None
         assert Path(tmpfile.name).stat().st_size > 0  # check that outgrid exists
-        temp_grid = load_dataarray(tmpfile.name)
+        temp_grid = read(tmpfile.name, kind="grid")
         xr.testing.assert_allclose(a=temp_grid, b=expected_grid)
 
 

pygmt/tests/test_grdclip.py

@@ -6,7 +6,7 @@
 
 import pytest
 import xarray as xr
-from pygmt import grdclip, load_dataarray
+from pygmt import grdclip, read
 from pygmt.helpers import GMTTempFile
 from pygmt.helpers.testing import load_static_earth_relief
 
@@ -49,7 +49,7 @@ def test_grdclip_outgrid(grid, expected_grid):
         )
         assert result is None  # return value is None
         assert Path(tmpfile.name).stat().st_size > 0  # check that outgrid exists
-        temp_grid = load_dataarray(tmpfile.name)
+        temp_grid = read(tmpfile.name, kind="grid")
         assert temp_grid.dims == ("lat", "lon")
         assert temp_grid.gmt.gtype == 1  # Geographic grid
         assert temp_grid.gmt.registration == 1  # Pixel registration

pygmt/tests/test_grdcut.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pytest
 import xarray as xr
-from pygmt import grdcut, load_dataarray
+from pygmt import grdcut, read
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import GMTTempFile
 from pygmt.helpers.testing import load_static_earth_relief
@@ -50,7 +50,7 @@ def test_grdcut_dataarray_in_file_out(grid, expected_grid, region):
     with GMTTempFile(suffix=".nc") as tmpfile:
         result = grdcut(grid, outgrid=tmpfile.name, region=region)
         assert result is None  # grdcut returns None if output to a file
-        temp_grid = load_dataarray(tmpfile.name)
+        temp_grid = read(tmpfile.name, kind="grid")
         xr.testing.assert_allclose(a=temp_grid, b=expected_grid)
 
 

pygmt/tests/test_grdfill.py

@@ -7,7 +7,7 @@
 import numpy as np
 import pytest
 import xarray as xr
-from pygmt import grdfill, load_dataarray
+from pygmt import grdfill, read
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import GMTTempFile
 from pygmt.helpers.testing import load_static_earth_relief
@@ -109,7 +109,7 @@ def test_grdfill_file_out(grid, expected_grid):
         result = grdfill(grid=grid, mode="c20", outgrid=tmpfile.name)
         assert result is None  # return value is None
         assert Path(tmpfile.name).stat().st_size > 0  # check that outfile exists
-        temp_grid = load_dataarray(tmpfile.name)
+        temp_grid = read(tmpfile.name, kind="grid")
         xr.testing.assert_allclose(a=temp_grid, b=expected_grid)
 
 

pygmt/tests/test_grdfilter.py

@@ -7,7 +7,7 @@
 import numpy as np
 import pytest
 import xarray as xr
-from pygmt import grdfilter, load_dataarray
+from pygmt import grdfilter, read
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import GMTTempFile
 from pygmt.helpers.testing import load_static_earth_relief
@@ -70,7 +70,7 @@ def test_grdfilter_dataarray_in_file_out(grid, expected_grid):
         )
         assert result is None  # return value is None
         assert Path(tmpfile.name).stat().st_size > 0  # check that outgrid exists
-        temp_grid = load_dataarray(tmpfile.name)
+        temp_grid = read(tmpfile.name, kind="grid")
         xr.testing.assert_allclose(a=temp_grid, b=expected_grid)
 
 

pygmt/tests/test_grdgradient.py

@@ -6,7 +6,7 @@
 
 import pytest
 import xarray as xr
-from pygmt import grdgradient, load_dataarray
+from pygmt import grdgradient, read
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import GMTTempFile
 from pygmt.helpers.testing import load_static_earth_relief
@@ -49,7 +49,7 @@ def test_grdgradient_outgrid(grid, expected_grid):
         )
         assert result is None  # return value is None
         assert Path(tmpfile.name).stat().st_size > 0  # check that outgrid exists
-        temp_grid = load_dataarray(tmpfile.name)
+        temp_grid = read(tmpfile.name, kind="grid")
         xr.testing.assert_allclose(a=temp_grid, b=expected_grid)