Skip to content

Commit

Permalink
Support mutli-segment stringio input
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman committed Sep 13, 2024
1 parent dece315 commit 486fce7
Showing 1 changed file with 25 additions and 16 deletions.
41 changes: 25 additions & 16 deletions pygmt/clib/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand Down Expand Up @@ -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()

Check warning on line 1667 in pygmt/clib/session.py

View check run for this annotation

Codecov / codecov/patch

pygmt/clib/session.py#L1664-L1667

Added lines #L1664 - L1667 were not covered by tests
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(
Expand All @@ -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()

Check warning on line 1691 in pygmt/clib/session.py

View check run for this annotation

Codecov / codecov/patch

pygmt/clib/session.py#L1691

Added line #L1691 was not covered by tests
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,
Expand Down

0 comments on commit 486fce7

Please sign in to comment.