diff --git a/doc/api/index.rst b/doc/api/index.rst index 60a1240150e..0349fae1b3e 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -293,6 +293,7 @@ conversion of Python variables to GMT virtual files: clib.Session.virtualfile_from_matrix clib.Session.virtualfile_from_vectors clib.Session.virtualfile_from_grid + clib.Session.virtualfile_out Low level access (these are mostly used by the :mod:`pygmt.clib` package): diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 3aa42c885df..652c3c65d9f 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -1609,6 +1609,69 @@ def virtualfile_from_data( # noqa: PLR0912 return file_context + @contextlib.contextmanager + def virtualfile_out( + self, kind: Literal["dataset", "grid"] = "dataset", fname: str | None = None + ): + r""" + Create a virtual file or an actual file for storing output data. + + If ``fname`` is not given, a virtual file will be created to store the output + data into a GMT data container and the function yields the name of the virutal + file. Otherwise, the output data will be written into the specified file and the + function simply yields the actual file name. + + Parameters + ---------- + kind + The data kind of the virtual file to create. Valid values are ``"dataset"`` + and ``"grid"``. Ignored if ``fname`` is specified. + fname + The name of the actual file to write the output data. No virtual file will + be created. + + Yields + ------ + vfile : str + Name of the virtual file or the actual file. + + Examples + -------- + >>> from pygmt.clib import Session + >>> from pygmt.datatypes import _GMT_DATASET + >>> from pygmt.helpers import GMTTempFile + >>> + >>> with GMTTempFile(suffix=".txt") as tmpfile: + ... with open(tmpfile.name, mode="w") as fp: + ... print("1.0 2.0 3.0 TEXT", file=fp) + ... + ... # Create a virtual file for storing the output table. + ... with Session() as lib: + ... with lib.virtualfile_out(kind="dataset") as vouttbl: + ... lib.call_module("read", f"{tmpfile.name} {vouttbl} -Td") + ... ds = lib.read_virtualfile(vouttbl, kind="dataset") + ... assert isinstance(ds.contents, _GMT_DATASET) + ... + ... # Write data to an actual file without creating a virtual file. + ... with Session() as lib: + ... with lib.virtualfile_out(fname=tmpfile.name) as vouttbl: + ... assert vouttbl == tmpfile.name + ... lib.call_module("read", f"{tmpfile.name} {vouttbl} -Td") + ... with open(vouttbl, mode="r") as fp: + ... line = fp.readline() + ... assert line == "1\t2\t3\tTEXT\n" + """ + if fname is not None: # Yield the actual file name. + yield fname + else: # Create a virtual file for storing the output data. + # Determine the family and geometry from kind + family, geometry = { + "dataset": ("GMT_IS_DATASET", "GMT_IS_PLP"), + "grid": ("GMT_IS_GRID", "GMT_IS_SURFACE"), + }[kind] + with self.open_virtualfile(family, geometry, "GMT_OUT", None) as vfile: + yield vfile + def read_virtualfile( self, vfname: str, kind: Literal["dataset", "grid", None] = None ): @@ -1633,27 +1696,23 @@ def read_virtualfile( ... with GMTTempFile(suffix=".txt") as tmpfile: ... with open(tmpfile.name, mode="w") as fp: ... print("1.0 2.0 3.0 TEXT", file=fp) - ... with lib.open_virtualfile( - ... "GMT_IS_DATASET", "GMT_IS_PLP", "GMT_OUT", None - ... ) as vfile: - ... lib.call_module("read", f"{tmpfile.name} {vfile} -Td") + ... with lib.virtualfile_out(kind="dataset") as vouttbl: + ... lib.call_module("read", f"{tmpfile.name} {vouttbl} -Td") ... # Read the virtual file as a void pointer - ... void_pointer = lib.read_virtualfile(vfile) + ... void_pointer = lib.read_virtualfile(vouttbl) ... assert isinstance(void_pointer, int) # void pointer is an int ... # Read the virtual file as a dataset - ... data_pointer = lib.read_virtualfile(vfile, kind="dataset") + ... data_pointer = lib.read_virtualfile(vouttbl, kind="dataset") ... assert isinstance(data_pointer, ctp.POINTER(_GMT_DATASET)) >>> >>> # Read grid from a virtual file >>> with Session() as lib: - ... with lib.open_virtualfile( - ... "GMT_IS_GRID", "GMT_IS_SURFACE", "GMT_OUT", None - ... ) as vfile: - ... lib.call_module("read", f"@earth_relief_01d_g {vfile} -Tg") + ... with lib.virtualfile_out(kind="grid") as voutgrd: + ... lib.call_module("read", f"@earth_relief_01d_g {voutgrd} -Tg") ... # Read the virtual file as a void pointer - ... void_pointer = lib.read_virtualfile(vfile) + ... void_pointer = lib.read_virtualfile(voutgrd) ... assert isinstance(void_pointer, int) # void pointer is an int - ... data_pointer = lib.read_virtualfile(vfile, kind="grid") + ... data_pointer = lib.read_virtualfile(voutgrd, kind="grid") ... assert isinstance(data_pointer, ctp.POINTER(_GMT_GRID)) Returns