Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

clib: Add the virtualfile_out method for creating output virtualfile #3057

Merged
merged 12 commits into from
Feb 28, 2024
1 change: 1 addition & 0 deletions doc/api/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
83 changes: 71 additions & 12 deletions pygmt/clib/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
):
Expand All @@ -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
Expand Down
Loading