Merge branch 'api/virtualfile-from-stringio' into legend/stringio

Author: seisman

doc/api/index.rst

@@ -317,5 +317,6 @@ Low level access (these are mostly used by the :mod:`pygmt.clib` package):
     clib.Session.get_libgmt_func
     clib.Session.virtualfile_from_data
     clib.Session.virtualfile_from_grid
+    clib.Session.virtualfile_from_stringio
     clib.Session.virtualfile_from_matrix
     clib.Session.virtualfile_from_vectors

pygmt/clib/session.py

@@ -7,6 +7,7 @@
 
 import contextlib
 import ctypes as ctp
+import io
 import pathlib
 import sys
 import warnings
@@ -60,6 +61,7 @@
     "GMT_IS_PLP",  # items could be any one of POINT, LINE, or POLY
     "GMT_IS_SURFACE",  # items are 2-D grid
     "GMT_IS_VOLUME",  # items are 3-D grid
+    "GMT_IS_TEXT",  # Text strings which triggers ASCII text reading
 ]
 
 METHODS = [
@@ -70,6 +72,11 @@
 DIRECTIONS = ["GMT_IN", "GMT_OUT"]
 
 MODES = ["GMT_CONTAINER_ONLY", "GMT_IS_OUTPUT"]
+MODE_MODIFIERS = [
+    "GMT_GRID_IS_CARTESIAN",
+    "GMT_GRID_IS_GEO",
+    "GMT_WITH_STRINGS",
+]
 
 REGISTRATIONS = ["GMT_GRID_PIXEL_REG", "GMT_GRID_NODE_REG"]
 
@@ -728,7 +735,7 @@ def create_data(
         mode_int = self._parse_constant(
             mode,
             valid=MODES,
-            valid_modifiers=["GMT_GRID_IS_CARTESIAN", "GMT_GRID_IS_GEO"],
+            valid_modifiers=MODE_MODIFIERS,
         )
         geometry_int = self._parse_constant(geometry, valid=GEOMETRIES)
         registration_int = self._parse_constant(registration, valid=REGISTRATIONS)
@@ -1603,6 +1610,83 @@ def virtualfile_from_grid(self, grid):
         with self.open_virtualfile(*args) as vfile:
             yield vfile
 
+    @contextlib.contextmanager
+    def virtualfile_from_stringio(self, stringio: io.StringIO):
+        r"""
+        Store a :class:`io.StringIO` object in a virtual file.
+
+        Store the contents of a :class:`io.StringIO` object in a GMT_DATASET container
+        and create a virtual file to pass to a GMT module.
+
+        Parameters
+        ----------
+        stringio
+            The :class:`io.StringIO` object containing the data to be stored in the
+            virtual file.
+
+        Yields
+        ------
+        fname
+            The name of the virtual file.
+
+        Examples
+        --------
+        >>> import io
+        >>> from pygmt.clib import Session
+        >>> stringio = 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"
+        ... )
+        >>> with Session() as lib:
+        ...     with lib.virtualfile_from_stringio(stringio) as fin:
+        ...         lib.virtualfile_to_dataset(vfname=fin, output_type="pandas")
+                                                     0
+        0                                 H 24p Legend
+        1                                          N 2
+        2  S 0.1i c 0.15i p300/12 0.25p 0.3i My circle
+        """
+        # Parse the strings in the io.StringIO object.
+        # For simplicity, we make a few assumptions.
+        # - "#" indicates a comment line
+        # - ">" indicates a segment header
+        # - Only one table and one segment
+        header = None
+        string_arrays = []
+        for line in stringio.getvalue().splitlines():
+            if line.startswith("#"):  # Skip comments
+                continue
+            if line.startswith(">"):  # Segment header
+                if header is not None:  # Only one segment is allowed now.
+                    raise GMTInvalidInput("Only one segment is allowed.")
+                header = line
+                continue
+            string_arrays.append(line)
+        # Only one table and one segment. No numeric data, so n_columns is 0.
+        n_tables, n_segments, n_rows, n_columns = 1, 1, len(string_arrays), 0
+
+        family, geometry = "GMT_IS_DATASET", "GMT_IS_TEXT"
+        dataset = self.create_data(
+            family,
+            geometry,
+            mode="GMT_CONTAINER_ONLY|GMT_WITH_STRINGS",
+            dim=[n_tables, n_segments, n_rows, n_columns],
+        )
+        dataset = ctp.cast(dataset, ctp.POINTER(_GMT_DATASET))
+        # Assign the strings to the segment
+        seg = dataset.contents.table[0].contents.segment[0].contents
+        if header is not None:
+            seg.header = header.encode()
+        seg.text = strings_to_ctypes_array(string_arrays)
+
+        with self.open_virtualfile(family, geometry, "GMT_IN", dataset) as vfile:
+            try:
+                yield vfile
+            finally:
+                # Must set the text to None to avoid double freeing the memory
+                seg.text = None
+
     def virtualfile_in(  # noqa: PLR0912
         self,
         check_kind=None,
@@ -1696,6 +1780,7 @@ def virtualfile_in(  # noqa: PLR0912
             "geojson": tempfile_from_geojson,
             "grid": self.virtualfile_from_grid,
             "image": tempfile_from_image,
+            "stringio": self.virtualfile_from_stringio,
             # Note: virtualfile_from_matrix is not used because a matrix can be
             # converted to vectors instead, and using vectors allows for better
             # handling of string type inputs (e.g. for datetime data types)
@@ -1704,7 +1789,7 @@ def virtualfile_in(  # noqa: PLR0912
         }[kind]
 
         # Ensure the data is an iterable (Python list or tuple)
-        if kind in {"geojson", "grid", "image", "file", "arg"}:
+        if kind in {"geojson", "grid", "image", "file", "arg", "stringio"}:
             if kind == "image" and data.dtype != "uint8":
                 msg = (
                     f"Input image has dtype: {data.dtype} which is unsupported, "

pygmt/helpers/utils.py

@@ -2,6 +2,7 @@
 Utilities and common tasks for wrapping the GMT modules.
 """
 
+import io
 import os
 import pathlib
 import shutil
@@ -188,8 +189,10 @@ def _check_encoding(
 
 def data_kind(
     data: Any = None, required: bool = True
-) -> Literal["arg", "file", "geojson", "grid", "image", "matrix", "vectors"]:
-    """
+) -> Literal[
+    "arg", "file", "geojson", "grid", "image", "matrix", "stringio", "vectors"
+]:
+    r"""
     Check the kind of data that is provided to a module.
 
     The ``data`` argument can be in any type, but only following types are supported:
@@ -222,6 +225,7 @@ def data_kind(
     >>> import numpy as np
     >>> import xarray as xr
     >>> import pathlib
+    >>> import io
     >>> data_kind(data=None)
     'vectors'
     >>> data_kind(data=np.arange(10).reshape((5, 2)))
@@ -240,8 +244,12 @@ def data_kind(
     'grid'
     >>> data_kind(data=xr.DataArray(np.random.rand(3, 4, 5)))
     'image'
+    >>> data_kind(data=io.StringIO("TEXT1\nTEXT23\n"))
+    'stringio'
     """
-    kind: Literal["arg", "file", "geojson", "grid", "image", "matrix", "vectors"]
+    kind: Literal[
+        "arg", "file", "geojson", "grid", "image", "matrix", "stringio", "vectors"
+    ]
     if isinstance(data, str | pathlib.PurePath) or (
         isinstance(data, list | tuple)
         and all(isinstance(_file, str | pathlib.PurePath) for _file in data)
@@ -250,6 +258,8 @@ def data_kind(
         kind = "file"
     elif isinstance(data, bool | int | float) or (data is None and not required):
         kind = "arg"
+    elif isinstance(data, io.StringIO):
+        kind = "stringio"
     elif isinstance(data, xr.DataArray):
         kind = "image" if len(data.dims) == 3 else "grid"
     elif hasattr(data, "__geo_interface__"):

pygmt/src/legend.py

@@ -97,4 +97,4 @@ def legend(
 
     with Session() as lib:
         with lib.virtualfile_in(data=spec, required_data=False) as vintbl:
-            lib.call_module(module="legend", args=build_arg_list(kwargs, infile=vintbl))
+            lib.call_module(module="legend", args=build_arg_list(kwargs, infile=vintbl))