Skip to content

Commit 486fce7

Browse files
committed
Support mutli-segment stringio input
1 parent dece315 commit 486fce7

File tree

1 file changed

+25
-16
lines changed

1 file changed

+25
-16
lines changed

pygmt/clib/session.py

+25-16
Original file line numberDiff line numberDiff line change
@@ -1622,7 +1622,7 @@ def virtualfile_from_stringio(self, stringio: io.StringIO):
16221622
16231623
- ``"#"`` indicates a comment line.
16241624
- ``">"`` indicates a segment header.
1625-
- The object only contains one table and one segment.
1625+
- The object only contains one table.
16261626
16271627
Parameters
16281628
----------
@@ -1655,19 +1655,26 @@ def virtualfile_from_stringio(self, stringio: io.StringIO):
16551655
2 S 0.1i c 0.15i p300/12 0.25p 0.3i My circle
16561656
"""
16571657
# Parse the strings in the io.StringIO object.
1658-
header = None
1659-
string_arrays = []
1658+
segments = []
1659+
current_segment = {"header": "", "data": []}
16601660
for line in stringio.getvalue().splitlines():
16611661
if line.startswith("#"): # Skip comments
16621662
continue
16631663
if line.startswith(">"): # Segment header
1664-
if header is not None: # Only one segment is allowed now.
1665-
raise GMTInvalidInput("Only one segment is allowed.")
1666-
header = line
1667-
continue
1668-
string_arrays.append(line)
1669-
# Only one table and one segment. No numeric data, so n_columns is 0.
1670-
n_tables, n_segments, n_rows, n_columns = 1, 1, len(string_arrays), 0
1664+
if current_segment["data"]: # If we have data, start a new segment
1665+
segments.append(current_segment)
1666+
current_segment = {"header": "", "data": []}
1667+
current_segment["header"] = line.strip(">").strip()
1668+
else:
1669+
current_segment["data"].append(line)
1670+
if current_segment["data"]: # Add the last segment if it has data
1671+
segments.append(current_segment)
1672+
1673+
# One table with one or more segments. No numeric data, so n_columns is 0.
1674+
n_tables = 1
1675+
n_segments = len(segments)
1676+
n_rows = sum(len(segment["data"]) for segment in segments)
1677+
n_columns = 0
16711678

16721679
family, geometry = "GMT_IS_DATASET", "GMT_IS_TEXT"
16731680
dataset = self.create_data(
@@ -1677,18 +1684,20 @@ def virtualfile_from_stringio(self, stringio: io.StringIO):
16771684
dim=[n_tables, n_segments, n_rows, n_columns],
16781685
)
16791686
dataset = ctp.cast(dataset, ctp.POINTER(_GMT_DATASET))
1680-
# Assign the strings to the segment
1681-
seg = dataset.contents.table[0].contents.segment[0].contents
1682-
if header is not None:
1683-
seg.header = header.encode()
1684-
seg.text = strings_to_ctypes_array(string_arrays)
1687+
table = dataset.contents.table[0].contents
1688+
for i, segment in enumerate(segments):
1689+
seg = table.segment[i].contents
1690+
if segment["header"] != "":
1691+
seg.header = segment["header"].encode()
1692+
seg.text = strings_to_ctypes_array(segment["data"])
16851693

16861694
with self.open_virtualfile(family, geometry, "GMT_IN", dataset) as vfile:
16871695
try:
16881696
yield vfile
16891697
finally:
16901698
# Must set the text to None to avoid double freeing the memory
1691-
seg.text = None
1699+
for i in range(n_segments):
1700+
table.segment[i].contents.text = None
16921701

16931702
def virtualfile_in( # noqa: PLR0912
16941703
self,

0 commit comments

Comments
 (0)