Merge branch 'main' into refactor/data_kind

Author: seisman

.github/workflows/format-command.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       # Generate token from GenericMappingTools bot
-      - uses: actions/create-github-app-token@v1.10.4
+      - uses: actions/create-github-app-token@v1.11.0
         id: generate-token
         with:
           app-id: ${{ secrets.APP_ID }}

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 sys
 import warnings
 from collections.abc import Generator, Sequence
@@ -59,6 +60,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 = [
@@ -69,6 +71,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"]
 
@@ -727,7 +734,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)
@@ -1602,6 +1609,100 @@ 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.
+
+        For simplicity, currently we make following assumptions in the StringIO object
+
+        - ``"#"`` indicates a comment line.
+        - ``">"`` indicates a segment header.
+
+        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
+        >>> # A StringIO object containing legend specifications
+        >>> 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 io.StringIO object.
+        segments = []
+        current_segment = {"header": "", "data": []}
+        for line in stringio.getvalue().splitlines():
+            if line.startswith("#"):  # Skip comments
+                continue
+            if line.startswith(">"):  # Segment header
+                if current_segment["data"]:  # If we have data, start a new segment
+                    segments.append(current_segment)
+                    current_segment = {"header": "", "data": []}
+                current_segment["header"] = line.strip(">").lstrip()
+            else:
+                current_segment["data"].append(line)  # type: ignore[attr-defined]
+        if current_segment["data"]:  # Add the last segment if it has data
+            segments.append(current_segment)
+
+        # One table with one or more segments.
+        # n_rows is the maximum number of rows/records for all segments.
+        # n_columns is the number of numeric data columns, so it's 0 here.
+        n_tables = 1
+        n_segments = len(segments)
+        n_rows = max(len(segment["data"]) for segment in segments)
+        n_columns = 0
+
+        # Create the GMT_DATASET container
+        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))
+        table = dataset.contents.table[0].contents
+        for i, segment in enumerate(segments):
+            seg = table.segment[i].contents
+            if segment["header"]:
+                seg.header = segment["header"].encode()  # type: ignore[attr-defined]
+            seg.text = strings_to_ctypes_array(segment["data"])
+
+        with self.open_virtualfile(family, geometry, "GMT_IN", dataset) as vfile:
+            try:
+                yield vfile
+            finally:
+                # Must set the pointers to None to avoid double freeing the memory.
+                # Maybe upstream bug.
+                for i in range(n_segments):
+                    seg = table.segment[i].contents
+                    seg.header = None
+                    seg.text = None
+
     def virtualfile_in(  # noqa: PLR0912
         self,
         check_kind=None,
@@ -1697,12 +1798,13 @@ def virtualfile_in(  # noqa: PLR0912
             "grid": self.virtualfile_from_grid,
             "image": tempfile_from_image,
             "matrix": self.virtualfile_from_matrix,
+            "stringio": self.virtualfile_from_stringio,
             "vectors": self.virtualfile_from_vectors,
         }[kind]
 
         # Ensure the data is an iterable (Python list or tuple)
         match kind:
-            case "arg" | "file" | "geojson" | "grid":
+            case "arg" | "file" | "geojson" | "grid" | "stringio":
                 _data = (data,)
             case "image":
                 if data.dtype != "uint8":

pygmt/helpers/tempfile.py

@@ -12,17 +12,17 @@
 from packaging.version import Version
 
 
-def unique_name():
+def unique_name() -> str:
     """
     Generate a unique name.
 
-    Useful for generating unique names for figures (otherwise GMT will plot
-    everything on the same figure instead of creating a new one).
+    Useful for generating unique names for figures. Otherwise GMT will plot everything
+    on the same figure instead of creating a new one.
 
     Returns
     -------
-    name : str
-        A unique name generated by :func:`uuid.uuid4`
+    name
+        A unique name generated by :func:`uuid.uuid4`.
     """
     return uuid.uuid4().hex
 
@@ -31,15 +31,14 @@ class GMTTempFile:
     """
     Context manager for creating closed temporary files.
 
-    This class does not return a file-like object. So, you can't do
-    ``for line in GMTTempFile()``, for example, or pass it to things that
-    need file objects.
+    This class does not return a file-like object. So, you can't iterate over the object
+    like ``for line in GMTTempFile()``, or pass it to things that need a file object.
 
     Parameters
     ----------
-    prefix : str
+    prefix
         The temporary file name begins with the prefix.
-    suffix : str
+    suffix
         The temporary file name ends with the suffix.
 
     Examples
@@ -60,7 +59,10 @@ class GMTTempFile:
     [0. 0. 0.] [1. 1. 1.] [2. 2. 2.]
     """
 
-    def __init__(self, prefix="pygmt-", suffix=".txt"):
+    def __init__(self, prefix: str = "pygmt-", suffix: str = ".txt"):
+        """
+        Initialize the object.
+        """
         with NamedTemporaryFile(prefix=prefix, suffix=suffix, delete=False) as tmpfile:
             self.name = tmpfile.name
 
@@ -76,33 +78,33 @@ def __exit__(self, *args):
         """
         Path(self.name).unlink(missing_ok=True)
 
-    def read(self, keep_tabs=False):
+    def read(self, keep_tabs: bool = False) -> str:
         """
         Read the entire contents of the file as a Unicode string.
 
         Parameters
         ----------
-        keep_tabs : bool
+        keep_tabs
             If False, replace the tabs that GMT uses with spaces.
 
         Returns
         -------
-        content : str
+        content
             Content of the temporary file as a Unicode string.
         """
         content = Path(self.name).read_text(encoding="utf8")
         if not keep_tabs:
             content = content.replace("\t", " ")
         return content
 
-    def loadtxt(self, **kwargs):
+    def loadtxt(self, **kwargs) -> np.ndarray:
         """
         Load data from the temporary file using numpy.loadtxt.
 
         Parameters
         ----------
-        kwargs : dict
-            Any keyword arguments that can be passed to numpy.loadtxt.
+        kwargs
+            Any keyword arguments that can be passed to :func:`np.loadtxt`.
 
         Returns
         -------

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,9 +189,11 @@ def _check_encoding(
 
 
 def data_kind(  # noqa: PLR0911
-    data: Any, required: bool = True
-) -> Literal["none", "arg", "file", "geojson", "grid", "image", "matrix", "vectors"]:
-    """
+    data: Any = None, required: bool = True
+) -> Literal[
+    "none", "arg", "file", "geojson", "grid", "image", "matrix", "stringio", "vectors"
+]:
+    r"""
     Check the kind of data that is provided to a module.
 
     Recognized data kinds are:
@@ -230,6 +233,7 @@ def data_kind(  # noqa: PLR0911
     >>> import xarray as xr
     >>> import pandas as pd
     >>> import pathlib
+    >>> import io
     >>> data_kind(data=None)
     'none'
     >>> data_kind(data=None, required=False)
@@ -248,6 +252,8 @@ def data_kind(  # noqa: PLR0911
     'image'
     >>> data_kind(data=np.arange(10).reshape((5, 2)))
     'matrix'
+    >>> data_kind(data=io.StringIO("TEXT1\nTEXT23\n"))
+    'stringio'
     >>> data_kind(data=pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}))
     'vectors'
     >>> data_kind(data={"x": [1, 2], "y": [3, 4]})
@@ -268,6 +274,10 @@ def data_kind(  # noqa: PLR0911
     ):
         return "file"
 
+    # A StringIO object.
+    if isinstance(data, io.StringIO):
+        return "stringio"
+
     # An option argument, mainly for dealing with optional virtual files
     if isinstance(data, bool | int | float) or (data is None and not required):
         return "arg"