diff --git a/doc/api/index.rst b/doc/api/index.rst index 01fac7d7a89..4b720a12353 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -172,6 +172,7 @@ Input/output :toctree: generated load_dataarray + read GMT Defaults ------------ diff --git a/pygmt/__init__.py b/pygmt/__init__.py index f6d1040851f..dba5cf38d1c 100644 --- a/pygmt/__init__.py +++ b/pygmt/__init__.py @@ -54,6 +54,7 @@ makecpt, nearneighbor, project, + read, select, sph2grd, sphdistance, diff --git a/pygmt/helpers/testing.py b/pygmt/helpers/testing.py index 29dfd08df19..990ce5aa418 100644 --- a/pygmt/helpers/testing.py +++ b/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): diff --git a/pygmt/src/__init__.py b/pygmt/src/__init__.py index e4db7321963..01cea172417 100644 --- a/pygmt/src/__init__.py +++ b/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 diff --git a/pygmt/src/read.py b/pygmt/src/read.py new file mode 100644 index 00000000000..5cd390f04b1 --- /dev/null +++ b/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) + + + Read a grid into an :class:`xarray.DataArray` object: + >>> dataarray = read("@earth_relief_01d", kind="grid") + >>> type(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)