From 486fce7f40ef99790b45a18256406ae4f62fe374 Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Fri, 13 Sep 2024 14:23:11 +0800
Subject: [PATCH] Support mutli-segment stringio input

---
 pygmt/clib/session.py | 41 +++++++++++++++++++++++++----------------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py
index e9d56d4bb1a..7baeea6ba31 100644
--- a/pygmt/clib/session.py
+++ b/pygmt/clib/session.py
@@ -1622,7 +1622,7 @@ def virtualfile_from_stringio(self, stringio: io.StringIO):
 
         - ``"#"`` indicates a comment line.
         - ``">"`` indicates a segment header.
-        - The object only contains one table and one segment.
+        - The object only contains one table.
 
         Parameters
         ----------
@@ -1655,19 +1655,26 @@ def virtualfile_from_stringio(self, stringio: io.StringIO):
         2  S 0.1i c 0.15i p300/12 0.25p 0.3i My circle
         """
         # Parse the strings in the io.StringIO object.
-        header = None
-        string_arrays = []
+        segments = []
+        current_segment = {"header": "", "data": []}
         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
+                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(">").strip()
+            else:
+                current_segment["data"].append(line)
+        if current_segment["data"]:  # Add the last segment if it has data
+            segments.append(current_segment)
+
+        # One table with one or more segments. No numeric data, so n_columns is 0.
+        n_tables = 1
+        n_segments = len(segments)
+        n_rows = sum(len(segment["data"]) for segment in segments)
+        n_columns = 0
 
         family, geometry = "GMT_IS_DATASET", "GMT_IS_TEXT"
         dataset = self.create_data(
@@ -1677,18 +1684,20 @@ def virtualfile_from_stringio(self, stringio: io.StringIO):
             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)
+        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()
+            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 text to None to avoid double freeing the memory
-                seg.text = None
+                for i in range(n_segments):
+                    table.segment[i].contents.text = None
 
     def virtualfile_in(  # noqa: PLR0912
         self,