Skip to content

Commit

Permalink
clib.Session: Add the virtualfile_to_data method for creating virtual…
Browse files Browse the repository at this point in the history
… files for output
  • Loading branch information
seisman committed Feb 20, 2024
1 parent ab39616 commit 5d278fc
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 0 deletions.
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_to_data


Low level access (these are mostly used by the :mod:`pygmt.clib` package):
Expand Down
68 changes: 68 additions & 0 deletions pygmt/clib/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -1609,6 +1609,74 @@ def virtualfile_from_data( # noqa: PLR0912

return file_context

@contextlib.contextmanager
def virtualfile_to_data(
self, kind: Literal["dataset", "grid", None] = None, fname: str | None = None
):
"""
Create a virtual file for storing output data in a data container or yield the
actual file name.
Parameters
----------
kind
The kind of data container to create. Valid values are ``"dataset"`` and
``"grid"`` or ``None``. Ignored if ``fname`` is specified.
fname
If given, yield the actual file name instead of the virtual file name.
Yields
------
vfile : str
Name of the virtual file or the output file name.
Examples
--------
>>> from pathlib import Path
>>> from pygmt.clib import Session
>>> from pygmt.datatypes import _GMT_DATASET, _GMT_GRID
>>> from pygmt.helpers import GMTTempFile
>>>
>>> # Create a virtual file for storing the output table.
>>> 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 Session() as lib:
... with lib.virtualfile_to_data(kind="dataset") as vouttbl:
... lib.call_module("read", f"{tmpfile.name} {vouttbl} -Td")
... ds = lib.read_virtualfile(vouttbl, kind="dataset")
>>> isinstance(ds.contents, _GMT_DATASET)
True
>>>
>>> # Create a virtual file for storing the output grid.
>>> with Session() as lib:
... with lib.virtualfile_to_data(kind="grid") as voutgrd:
... lib.call_module("read", f"@earth_relief_01d_g {voutgrd} -Tg")
... outgrd = lib.read_virtualfile(voutgrd, kind="grid")
>>> isinstance(outgrd.contents, _GMT_GRID)
True
>>>
>>> # Write data to file without creating a virtual file
>>> with GMTTempFile(suffix=".nc") as tmpfile:
... with Session() as lib:
... with lib.virtualfile_to_data(fname=tmpfile.name) as voutgrd:
... lib.call_module("read", f"@earth_relief_01d_g {voutgrd} -Tg")
... assert voutgrd == tmpfile.name
... assert Path(voutgrd).stat().st_size > 0
"""
# If fname is given, yield the output file name.
if fname is not None:
yield fname
# Otherwise, create a virtual file for storing the output data.
else:
# Determine the family and geometry from kind
family, geometry = {
"grid": ("GMT_IS_GRID", "GMT_IS_SURFACE"),
"dataset": ("GMT_IS_DATASET", "GMT_IS_PLP"),
}[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 Down

0 comments on commit 5d278fc

Please sign in to comment.