Split tests related to virtualfiles into multiple test files

Author: seisman

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_open_virtualfile.py

@@ -0,0 +1,122 @@
+"""
+Test the Session.open_virtualfile method.
+"""
+
+from importlib.util import find_spec
+from pathlib import Path
+
+import numpy as np
+import pytest
+from pygmt import clib
+from pygmt.exceptions import GMTCLibError, GMTInvalidInput
+from pygmt.helpers import GMTTempFile
+from pygmt.tests.test_clib import mock
+
+POINTS_DATA = Path(__file__).parent / "data" / "points.txt"
+
+
+@pytest.fixture(scope="module", name="data")
+def fixture_data():
+    """
+    Load the point data from the test file.
+    """
+    return np.loadtxt(POINTS_DATA)
+
+
+@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.fixture(scope="module", name="dtypes_pandas")
+def fixture_dtypes_pandas(dtypes):
+    """
+    List of supported pandas dtypes.
+    """
+    dtypes_pandas = dtypes.copy()
+
+    if find_spec("pyarrow") is not None:
+        dtypes_pandas.extend([f"{dtype}[pyarrow]" for dtype in dtypes_pandas])
+
+    return tuple(dtypes_pandas)
+
+
+@pytest.mark.benchmark
+def test_open_virtualfile(dtypes):
+    """
+    Test passing in data via a virtual file with a Dataset.
+    """
+    shape = (5, 3)
+    for dtype in dtypes:
+        with clib.Session() as lib:
+            family = "GMT_IS_DATASET|GMT_VIA_MATRIX"
+            geometry = "GMT_IS_POINT"
+            dataset = lib.create_data(
+                family=family,
+                geometry=geometry,
+                mode="GMT_CONTAINER_ONLY",
+                dim=[shape[1], shape[0], 1, 0],  # columns, rows, layers, dtype
+            )
+            data = np.arange(shape[0] * shape[1], dtype=dtype).reshape(shape)
+            lib.put_matrix(dataset, matrix=data)
+            # Add the dataset to a virtual file and pass it along to gmt info
+            vfargs = (family, geometry, "GMT_IN|GMT_IS_REFERENCE", dataset)
+            with lib.open_virtualfile(*vfargs) 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_open_virtualfile_fails():
+    """
+    Check that opening and closing virtual files raises an exception for non- zero
+    return codes.
+    """
+    vfargs = (
+        "GMT_IS_DATASET|GMT_VIA_MATRIX",
+        "GMT_IS_POINT",
+        "GMT_IN|GMT_IS_REFERENCE",
+        None,
+    )
+
+    # Mock Open_VirtualFile to test the status check when entering the context.
+    # If the exception is raised, the code won't get to the closing of the
+    # virtual file.
+    with clib.Session() as lib, mock(lib, "GMT_Open_VirtualFile", returns=1):
+        with pytest.raises(GMTCLibError):
+            with lib.open_virtualfile(*vfargs):
+                pass
+
+    # Test the status check when closing the virtual file
+    # Mock the opening to return 0 (success) so that we don't open a file that
+    # we won't close later.
+    with (
+        clib.Session() as lib,
+        mock(lib, "GMT_Open_VirtualFile", returns=0),
+        mock(lib, "GMT_Close_VirtualFile", returns=1),
+    ):
+        with pytest.raises(GMTCLibError):
+            with lib.open_virtualfile(*vfargs):
+                pass
+
+
+def test_open_virtualfile_bad_direction():
+    """
+    Test passing an invalid direction argument.
+    """
+    with clib.Session() as lib:
+        vfargs = (
+            "GMT_IS_DATASET|GMT_VIA_MATRIX",
+            "GMT_IS_POINT",
+            "GMT_IS_GRID",  # The invalid direction argument
+            0,
+        )
+        with pytest.raises(GMTInvalidInput):
+            with lib.open_virtualfile(*vfargs):
+                pass

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