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

Author: seisman

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/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.
 
     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,76 @@
+"""
+Read data from files
+"""
+
+from typing import Literal
+
+import pandas as pd
+import xarray as xr
+from pygmt.clib import Session
+
+
+def read(
+    file, kind: Literal["dataset", "grid", "image"], **kwargs
+) -> pd.DataFrame | xr.DataArray:
+    """
+    Read a dataset, grid, or image from a file and return the appropriate object.
+
+    For datasets, it returns a :class:`pandas.DataFrame`. For grids and images,
+    it returns a :class:`xarray.DataArray`.
+
+    Parameters
+    ----------
+    file
+        The file name to read.
+    kind
+        The kind of data to read. Valid values are ``"dataset"``, ``"grid"``,
+        and ``"image"``.
+
+    For datasets, the following keyword arguments are supported:
+
+    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
+    -------
+    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'>
+    """
+    code = {"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, f"-T{code}"])
+
+        match kind:
+            case "dataset":
+                return lib.virtualfile_to_dataset(vfname=voutfile, **kwargs)
+            case "grid" | "image":
+                return lib.virtualfile_to_raster(vfname=voutfile, kind=kind)