Merge branch 'main' into refactor/data_kind

Author: seisman

pygmt/clib/session.py

@@ -1203,41 +1203,46 @@ def write_data(self, family, geometry, mode, wesn, output, data):
             raise GMTCLibError(f"Failed to write dataset to '{output}'")
 
     @contextlib.contextmanager
-    def open_virtualfile(self, family, geometry, direction, data):
+    def open_virtualfile(
+        self,
+        family: str,
+        geometry: str,
+        direction: str,
+        data: ctp.c_void_p | None,
+    ) -> Generator[str, None, None]:
         """
-        Open a GMT virtual file to pass data to and from a module.
+        Open a GMT virtual file associated with a data object for reading or writing.
 
-        GMT uses a virtual file scheme to pass in data or get data from API
-        modules. Use it to pass in your GMT data structure (created using
-        :meth:`pygmt.clib.Session.create_data`) to a module that expects an
-        input file, or get the output from a module that writes to a file.
+        GMT uses a virtual file scheme to pass in data or get data from API modules. Use
+        it to pass in your GMT data structure (created using
+        :meth:`pygmt.clib.Session.create_data`) to a module that expects an input file,
+        or get the output from a module that writes to a file.
 
-        Use in a ``with`` block. Will automatically close the virtual file when
-        leaving the ``with`` block. Because of this, no wrapper for
-        ``GMT_Close_VirtualFile`` is provided.
+        Use in a ``with`` block. Will automatically close the virtual file when leaving
+        the ``with`` block. Because of this, no wrapper for ``GMT_Close_VirtualFile``
+        is provided.
 
         Parameters
         ----------
-        family : str
-            A valid GMT data family name (e.g., ``"GMT_IS_DATASET"``). Should
-            be the same as the one you used to create your data structure.
-        geometry : str
-            A valid GMT data geometry name (e.g., ``"GMT_IS_POINT"``). Should
-            be the same as the one you used to create your data structure.
-        direction : str
-            Either ``"GMT_IN"`` or ``"GMT_OUT"`` to indicate if passing data to
-            GMT or getting it out of GMT, respectively.
-            By default, GMT can modify the data you pass in. Add modifier
-            ``"GMT_IS_REFERENCE"`` to tell GMT the data are read-only, or
-            ``"GMT_IS_DUPLICATE"`` to tell GMT to duplicate the data.
-        data : int or None
-            The ctypes void pointer to your GMT data structure. For output
-            (i.e., ``direction="GMT_OUT"``), it can be ``None`` to have GMT
-            automatically allocate the output GMT data structure.
+        family
+            A valid GMT data family name (e.g., ``"GMT_IS_DATASET"``). Should be the
+            same as the one you used to create your data structure.
+        geometry
+            A valid GMT data geometry name (e.g., ``"GMT_IS_POINT"``). Should be the
+            same as the one you used to create your data structure.
+        direction
+            Either ``"GMT_IN"`` or ``"GMT_OUT"`` to indicate if passing data to GMT or
+            getting it out of GMT, respectively. By default, GMT can modify the data you
+            pass in. Add modifier ``"GMT_IS_REFERENCE"`` to tell GMT the data are
+            read-only, or ``"GMT_IS_DUPLICATE"`` to tell GMT to duplicate the data.
+        data
+            The ctypes void pointer to the GMT data structure. For output (i.e.,
+            ``direction="GMT_OUT"``), it can be ``None`` to have GMT automatically
+            allocate the output GMT data structure.
 
         Yields
         ------
-        vfname : str
+        vfname
             The name of the virtual file that you can pass to a GMT module.
 
         Examples
@@ -1270,19 +1275,19 @@ def open_virtualfile(self, family, geometry, direction, data):
         c_open_virtualfile = self.get_libgmt_func(
             "GMT_Open_VirtualFile",
             argtypes=[
-                ctp.c_void_p,
-                ctp.c_uint,
-                ctp.c_uint,
-                ctp.c_uint,
-                ctp.c_void_p,
-                ctp.c_char_p,
+                ctp.c_void_p,  # V_API
+                ctp.c_uint,  # family
+                ctp.c_uint,  # geometry
+                ctp.c_uint,  # direction
+                ctp.c_void_p,  # data
+                ctp.c_char_p,  # name
             ],
             restype=ctp.c_int,
         )
 
         c_close_virtualfile = self.get_libgmt_func(
             "GMT_Close_VirtualFile",
-            argtypes=[ctp.c_void_p, ctp.c_char_p],
+            argtypes=[ctp.c_void_p, ctp.c_char_p],  # V_API, name
             restype=ctp.c_int,
         )
 
@@ -1297,19 +1302,24 @@ def open_virtualfile(self, family, geometry, direction, data):
             self.session_pointer, family_int, geometry_int, direction_int, data, buff
         )
         if status != 0:
-            raise GMTCLibError("Failed to create a virtual file.")
+            msg = (
+                f"Failed to create a virtual file with {family=}, {geometry=}, "
+                f"{direction=}."
+            )
+            raise GMTCLibError(msg)
 
         vfname = buff.value.decode()
         try:
             yield vfname
         finally:
             status = c_close_virtualfile(self.session_pointer, vfname.encode())
             if status != 0:
-                raise GMTCLibError(f"Failed to close virtual file '{vfname}'.")
+                msg = f"Failed to close virtual file '{vfname}'."
+                raise GMTCLibError(msg)
 
     def open_virtual_file(self, family, geometry, direction, data):
         """
-        Open a GMT virtual file to pass data to and from a module.
+        Open a GMT virtual file associated with a data object for reading or writing.
 
         .. deprecated: 0.11.0
 

pygmt/datasets/load_remote_dataset.py

@@ -441,7 +441,7 @@ def _load_remote_dataset(
             "returned unless only the pixel-registered grid is available."
         )
 
-    fname = f"@{prefix}_{resolution}_{registration[0]}"
+    fname = f"@{prefix}_{resolution}_{registration[0]}"  # type: ignore[index]
     if resinfo.tiled and region is None:
         raise GMTInvalidInput(
             f"'region' is required for {dataset.description} resolution '{resolution}'."

pygmt/tests/test_clib_inquire_virtualfile.py

@@ -0,0 +1,33 @@
+"""
+Test the Session.inquire_virtualfile method.
+"""
+
+from pygmt import clib
+
+
+def test_inquire_virtualfile():
+    """
+    Test that the inquire_virtualfile method returns the correct family.
+
+    Currently, only output virtual files are tested.
+    """
+    with clib.Session() as lib:
+        for family in [
+            "GMT_IS_DATASET",
+            "GMT_IS_DATASET|GMT_VIA_MATRIX",
+            "GMT_IS_DATASET|GMT_VIA_VECTOR",
+        ]:
+            with lib.open_virtualfile(
+                family, "GMT_IS_PLP", "GMT_OUT|GMT_IS_REFERENCE", None
+            ) as vfile:
+                assert lib.inquire_virtualfile(vfile) == lib["GMT_IS_DATASET"]
+
+        for family, geometry in [
+            ("GMT_IS_GRID", "GMT_IS_SURFACE"),
+            ("GMT_IS_IMAGE", "GMT_IS_SURFACE"),
+            ("GMT_IS_CUBE", "GMT_IS_VOLUME"),
+            ("GMT_IS_PALETTE", "GMT_IS_NONE"),
+            ("GMT_IS_POSTSCRIPT", "GMT_IS_NONE"),
+        ]:
+            with lib.open_virtualfile(family, geometry, "GMT_OUT", None) as vfile:
+                assert lib.inquire_virtualfile(vfile) == lib[family]

pygmt/tests/test_clib_virtualfile_from_matrix.py

@@ -0,0 +1,54 @@
+"""
+Test the Session.virtualfile_from_matrix method.
+"""
+
+import numpy as np
+import pytest
+from pygmt import clib
+from pygmt.helpers import GMTTempFile
+
+
+@pytest.fixture(scope="module", name="dtypes")
+def fixture_dtypes():
+    """
+    List of supported numpy dtypes.
+    """
+    return "int8 int16 int32 int64 uint8 uint16 uint32 uint64 float32 float64".split()
+
+
+@pytest.mark.benchmark
+def test_virtualfile_from_matrix(dtypes):
+    """
+    Test transforming a matrix to virtual file dataset.
+    """
+    shape = (7, 5)
+    for dtype in dtypes:
+        data = np.arange(shape[0] * shape[1], dtype=dtype).reshape(shape)
+        with clib.Session() as lib:
+            with lib.virtualfile_from_matrix(data) as vfile:
+                with GMTTempFile() as outfile:
+                    lib.call_module("info", [vfile, f"->{outfile.name}"])
+                    output = outfile.read(keep_tabs=True)
+            bounds = "\t".join([f"<{col.min():.0f}/{col.max():.0f}>" for col in data.T])
+            expected = f"<matrix memory>: N = {shape[0]}\t{bounds}\n"
+            assert output == expected
+
+
+def test_virtualfile_from_matrix_slice(dtypes):
+    """
+    Test transforming a slice of a larger array to virtual file dataset.
+    """
+    shape = (10, 6)
+    for dtype in dtypes:
+        full_data = np.arange(shape[0] * shape[1], dtype=dtype).reshape(shape)
+        rows = 5
+        cols = 3
+        data = full_data[:rows, :cols]
+        with clib.Session() as lib:
+            with lib.virtualfile_from_matrix(data) as vfile:
+                with GMTTempFile() as outfile:
+                    lib.call_module("info", [vfile, f"->{outfile.name}"])
+                    output = outfile.read(keep_tabs=True)
+            bounds = "\t".join([f"<{col.min():.0f}/{col.max():.0f}>" for col in data.T])
+            expected = f"<matrix memory>: N = {rows}\t{bounds}\n"
+            assert output == expected

pygmt/tests/test_clib_virtualfile_from_stringio.py

@@ -0,0 +1,106 @@
+"""
+Test the Session.virtualfile_from_stringio method.
+"""
+
+import io
+
+import numpy as np
+from pygmt import clib
+
+
+def _stringio_to_dataset(data: io.StringIO):
+    """
+    A helper function for check the virtualfile_from_stringio method.
+
+    The function does the following:
+
+    1. Creates a virtual file from the input StringIO object.
+    2. Pass the virtual file to the ``read`` module, which reads the virtual file and
+       writes it to another virtual file.
+    3. Reads the output virtual file as a GMT_DATASET object.
+    4. Extracts the header and the trailing text from the dataset and returns it as a
+       string.
+    """
+    with clib.Session() as lib:
+        with (
+            lib.virtualfile_from_stringio(data) as vintbl,
+            lib.virtualfile_out(kind="dataset") as vouttbl,
+        ):
+            lib.call_module("read", args=[vintbl, vouttbl, "-Td"])
+            ds = lib.read_virtualfile(vouttbl, kind="dataset").contents
+
+            output = []
+            table = ds.table[0].contents
+            for segment in table.segment[: table.n_segments]:
+                seg = segment.contents
+                output.append(f"> {seg.header.decode()}" if seg.header else ">")
+                output.extend(np.char.decode(seg.text[: seg.n_rows]))
+        return "\n".join(output) + "\n"
+
+
+def test_virtualfile_from_stringio():
+    """
+    Test the virtualfile_from_stringio method.
+    """
+    data = io.StringIO(
+        "# Comment\n"
+        "H 24p Legend\n"
+        "N 2\n"
+        "S 0.1i c 0.15i p300/12 0.25p 0.3i My circle\n"
+    )
+    expected = (
+        ">\n" "H 24p Legend\n" "N 2\n" "S 0.1i c 0.15i p300/12 0.25p 0.3i My circle\n"
+    )
+    assert _stringio_to_dataset(data) == expected
+
+
+def test_one_segment():
+    """
+    Test the virtualfile_from_stringio method with one segment.
+    """
+    data = io.StringIO(
+        "# Comment\n"
+        "> Segment 1\n"
+        "1 2 3 ABC\n"
+        "4 5 DE\n"
+        "6 7 8   9  FGHIJK LMN OPQ\n"
+        "RSTUVWXYZ\n"
+    )
+    expected = (
+        "> Segment 1\n"
+        "1 2 3 ABC\n"
+        "4 5 DE\n"
+        "6 7 8   9  FGHIJK LMN OPQ\n"
+        "RSTUVWXYZ\n"
+    )
+    assert _stringio_to_dataset(data) == expected
+
+
+def test_multiple_segments():
+    """
+    Test the virtualfile_from_stringio method with multiple segments.
+    """
+    data = io.StringIO(
+        "# Comment line 1\n"
+        "# Comment line 2\n"
+        "> Segment 1\n"
+        "1 2 3 ABC\n"
+        "4 5 DE\n"
+        "6 7 8   9  FG\n"
+        "# Comment line 3\n"
+        "> Segment 2\n"
+        "1 2 3 ABC\n"
+        "4 5 DE\n"
+        "6 7 8   9  FG\n"
+    )
+    expected = (
+        "> Segment 1\n"
+        "1 2 3 ABC\n"
+        "4 5 DE\n"
+        "6 7 8   9  FG\n"
+        "> Segment 2\n"
+        "1 2 3 ABC\n"
+        "4 5 DE\n"
+        "6 7 8   9  FG\n"
+    )
+    assert _stringio_to_dataset(data) == expected