Skip to content

Commit 2646d26

Browse files
seismanseisman
committed
Merge branch 'main' into refactor/data_kind
2 parents 5789aee + 890626d commit 2646d26

File tree

6 files changed

+242
-23
lines changed

6 files changed

+242
-23
lines changed

‎.github/workflows/format-command.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ jobs:
1111
runs-on: ubuntu-latest
1212
steps:
1313
# Generate token from GenericMappingTools bot
14-
- uses: actions/create-github-app-token@v1.10.4
14+
- uses: actions/create-github-app-token@v1.11.0
1515
id: generate-token
1616
with:
1717
app-id: ${{ secrets.APP_ID }}

‎doc/api/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,5 +317,6 @@ Low level access (these are mostly used by the :mod:`pygmt.clib` package):
317317
clib.Session.get_libgmt_func
318318
clib.Session.virtualfile_from_data
319319
clib.Session.virtualfile_from_grid
320+
clib.Session.virtualfile_from_stringio
320321
clib.Session.virtualfile_from_matrix
321322
clib.Session.virtualfile_from_vectors

‎pygmt/clib/session.py‎

Lines changed: 104 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import contextlib
99
import ctypes as ctp
10+
import io
1011
import sys
1112
import warnings
1213
from collections.abc import Generator, Sequence
@@ -59,6 +60,7 @@
5960
"GMT_IS_PLP", # items could be any one of POINT, LINE, or POLY
6061
"GMT_IS_SURFACE", # items are 2-D grid
6162
"GMT_IS_VOLUME", # items are 3-D grid
63+
"GMT_IS_TEXT", # Text strings which triggers ASCII text reading
6264
]
6365

6466
METHODS = [
@@ -69,6 +71,11 @@
6971
DIRECTIONS = ["GMT_IN", "GMT_OUT"]
7072

7173
MODES = ["GMT_CONTAINER_ONLY", "GMT_IS_OUTPUT"]
74+
MODE_MODIFIERS = [
75+
"GMT_GRID_IS_CARTESIAN",
76+
"GMT_GRID_IS_GEO",
77+
"GMT_WITH_STRINGS",
78+
]
7279

7380
REGISTRATIONS = ["GMT_GRID_PIXEL_REG", "GMT_GRID_NODE_REG"]
7481

@@ -727,7 +734,7 @@ def create_data(
727734
mode_int = self._parse_constant(
728735
mode,
729736
valid=MODES,
730-
valid_modifiers=["GMT_GRID_IS_CARTESIAN", "GMT_GRID_IS_GEO"],
737+
valid_modifiers=MODE_MODIFIERS,
731738
)
732739
geometry_int = self._parse_constant(geometry, valid=GEOMETRIES)
733740
registration_int = self._parse_constant(registration, valid=REGISTRATIONS)
@@ -1602,6 +1609,100 @@ def virtualfile_from_grid(self, grid):
16021609
with self.open_virtualfile(*args) as vfile:
16031610
yield vfile
16041611

1612+
@contextlib.contextmanager
1613+
def virtualfile_from_stringio(self, stringio: io.StringIO):
1614+
r"""
1615+
Store a :class:`io.StringIO` object in a virtual file.
1616+
1617+
Store the contents of a :class:`io.StringIO` object in a GMT_DATASET container
1618+
and create a virtual file to pass to a GMT module.
1619+
1620+
For simplicity, currently we make following assumptions in the StringIO object
1621+
1622+
- ``"#"`` indicates a comment line.
1623+
- ``">"`` indicates a segment header.
1624+
1625+
Parameters
1626+
----------
1627+
stringio
1628+
The :class:`io.StringIO` object containing the data to be stored in the
1629+
virtual file.
1630+
1631+
Yields
1632+
------
1633+
fname
1634+
The name of the virtual file.
1635+
1636+
Examples
1637+
--------
1638+
>>> import io
1639+
>>> from pygmt.clib import Session
1640+
>>> # A StringIO object containing legend specifications
1641+
>>> stringio = io.StringIO(
1642+
... "# Comment\n"
1643+
... "H 24p Legend\n"
1644+
... "N 2\n"
1645+
... "S 0.1i c 0.15i p300/12 0.25p 0.3i My circle\n"
1646+
... )
1647+
>>> with Session() as lib:
1648+
... with lib.virtualfile_from_stringio(stringio) as fin:
1649+
... lib.virtualfile_to_dataset(vfname=fin, output_type="pandas")
1650+
0
1651+
0 H 24p Legend
1652+
1 N 2
1653+
2 S 0.1i c 0.15i p300/12 0.25p 0.3i My circle
1654+
"""
1655+
# Parse the io.StringIO object.
1656+
segments = []
1657+
current_segment = {"header": "", "data": []}
1658+
for line in stringio.getvalue().splitlines():
1659+
if line.startswith("#"): # Skip comments
1660+
continue
1661+
if line.startswith(">"): # Segment header
1662+
if current_segment["data"]: # If we have data, start a new segment
1663+
segments.append(current_segment)
1664+
current_segment = {"header": "", "data": []}
1665+
current_segment["header"] = line.strip(">").lstrip()
1666+
else:
1667+
current_segment["data"].append(line) # type: ignore[attr-defined]
1668+
if current_segment["data"]: # Add the last segment if it has data
1669+
segments.append(current_segment)
1670+
1671+
# One table with one or more segments.
1672+
# n_rows is the maximum number of rows/records for all segments.
1673+
# n_columns is the number of numeric data columns, so it's 0 here.
1674+
n_tables = 1
1675+
n_segments = len(segments)
1676+
n_rows = max(len(segment["data"]) for segment in segments)
1677+
n_columns = 0
1678+
1679+
# Create the GMT_DATASET container
1680+
family, geometry = "GMT_IS_DATASET", "GMT_IS_TEXT"
1681+
dataset = self.create_data(
1682+
family,
1683+
geometry,
1684+
mode="GMT_CONTAINER_ONLY|GMT_WITH_STRINGS",
1685+
dim=[n_tables, n_segments, n_rows, n_columns],
1686+
)
1687+
dataset = ctp.cast(dataset, ctp.POINTER(_GMT_DATASET))
1688+
table = dataset.contents.table[0].contents
1689+
for i, segment in enumerate(segments):
1690+
seg = table.segment[i].contents
1691+
if segment["header"]:
1692+
seg.header = segment["header"].encode() # type: ignore[attr-defined]
1693+
seg.text = strings_to_ctypes_array(segment["data"])
1694+
1695+
with self.open_virtualfile(family, geometry, "GMT_IN", dataset) as vfile:
1696+
try:
1697+
yield vfile
1698+
finally:
1699+
# Must set the pointers to None to avoid double freeing the memory.
1700+
# Maybe upstream bug.
1701+
for i in range(n_segments):
1702+
seg = table.segment[i].contents
1703+
seg.header = None
1704+
seg.text = None
1705+
16051706
def virtualfile_in( # noqa: PLR0912
16061707
self,
16071708
check_kind=None,
@@ -1697,12 +1798,13 @@ def virtualfile_in( # noqa: PLR0912
16971798
"grid": self.virtualfile_from_grid,
16981799
"image": tempfile_from_image,
16991800
"matrix": self.virtualfile_from_matrix,
1801+
"stringio": self.virtualfile_from_stringio,
17001802
"vectors": self.virtualfile_from_vectors,
17011803
}[kind]
17021804

17031805
# Ensure the data is an iterable (Python list or tuple)
17041806
match kind:
1705-
case "arg" | "file" | "geojson" | "grid":
1807+
case "arg" | "file" | "geojson" | "grid" | "stringio":
17061808
_data = (data,)
17071809
case "image":
17081810
if data.dtype != "uint8":

‎pygmt/helpers/tempfile.py‎

Lines changed: 19 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,17 @@
1212
from packaging.version import Version
1313

1414

15-
def unique_name():
15+
def unique_name() -> str:
1616
"""
1717
Generate a unique name.
1818
19-
Useful for generating unique names for figures (otherwise GMT will plot
20-
everything on the same figure instead of creating a new one).
19+
Useful for generating unique names for figures. Otherwise GMT will plot everything
20+
on the same figure instead of creating a new one.
2121
2222
Returns
2323
-------
24-
name : str
25-
A unique name generated by :func:`uuid.uuid4`
24+
name
25+
A unique name generated by :func:`uuid.uuid4`.
2626
"""
2727
return uuid.uuid4().hex
2828

@@ -31,15 +31,14 @@ class GMTTempFile:
3131
"""
3232
Context manager for creating closed temporary files.
3333
34-
This class does not return a file-like object. So, you can't do
35-
``for line in GMTTempFile()``, for example, or pass it to things that
36-
need file objects.
34+
This class does not return a file-like object. So, you can't iterate over the object
35+
like ``for line in GMTTempFile()``, or pass it to things that need a file object.
3736
3837
Parameters
3938
----------
40-
prefix : str
39+
prefix
4140
The temporary file name begins with the prefix.
42-
suffix : str
41+
suffix
4342
The temporary file name ends with the suffix.
4443
4544
Examples
@@ -60,7 +59,10 @@ class GMTTempFile:
6059
[0. 0. 0.] [1. 1. 1.] [2. 2. 2.]
6160
"""
6261

63-
def __init__(self, prefix="pygmt-", suffix=".txt"):
62+
def __init__(self, prefix: str = "pygmt-", suffix: str = ".txt"):
63+
"""
64+
Initialize the object.
65+
"""
6466
with NamedTemporaryFile(prefix=prefix, suffix=suffix, delete=False) as tmpfile:
6567
self.name = tmpfile.name
6668

@@ -76,33 +78,33 @@ def __exit__(self, *args):
7678
"""
7779
Path(self.name).unlink(missing_ok=True)
7880

79-
def read(self, keep_tabs=False):
81+
def read(self, keep_tabs: bool = False) -> str:
8082
"""
8183
Read the entire contents of the file as a Unicode string.
8284
8385
Parameters
8486
----------
85-
keep_tabs : bool
87+
keep_tabs
8688
If False, replace the tabs that GMT uses with spaces.
8789
8890
Returns
8991
-------
90-
content : str
92+
content
9193
Content of the temporary file as a Unicode string.
9294
"""
9395
content = Path(self.name).read_text(encoding="utf8")
9496
if not keep_tabs:
9597
content = content.replace("\t", " ")
9698
return content
9799

98-
def loadtxt(self, **kwargs):
100+
def loadtxt(self, **kwargs) -> np.ndarray:
99101
"""
100102
Load data from the temporary file using numpy.loadtxt.
101103
102104
Parameters
103105
----------
104-
kwargs : dict
105-
Any keyword arguments that can be passed to numpy.loadtxt.
106+
kwargs
107+
Any keyword arguments that can be passed to :func:`np.loadtxt`.
106108
107109
Returns
108110
-------

‎pygmt/helpers/utils.py‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
Utilities and common tasks for wrapping the GMT modules.
33
"""
44

5+
import io
56
import os
67
import pathlib
78
import shutil
@@ -188,9 +189,11 @@ def _check_encoding(
188189

189190

190191
def data_kind( # noqa: PLR0911
191-
data: Any, required: bool = True
192-
) -> Literal["none", "arg", "file", "geojson", "grid", "image", "matrix", "vectors"]:
193-
"""
192+
data: Any = None, required: bool = True
193+
) -> Literal[
194+
"none", "arg", "file", "geojson", "grid", "image", "matrix", "stringio", "vectors"
195+
]:
196+
r"""
194197
Check the kind of data that is provided to a module.
195198
196199
Recognized data kinds are:
@@ -230,6 +233,7 @@ def data_kind( # noqa: PLR0911
230233
>>> import xarray as xr
231234
>>> import pandas as pd
232235
>>> import pathlib
236+
>>> import io
233237
>>> data_kind(data=None)
234238
'none'
235239
>>> data_kind(data=None, required=False)
@@ -248,6 +252,8 @@ def data_kind( # noqa: PLR0911
248252
'image'
249253
>>> data_kind(data=np.arange(10).reshape((5, 2)))
250254
'matrix'
255+
>>> data_kind(data=io.StringIO("TEXT1\nTEXT23\n"))
256+
'stringio'
251257
>>> data_kind(data=pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}))
252258
'vectors'
253259
>>> data_kind(data={"x": [1, 2], "y": [3, 4]})
@@ -268,6 +274,10 @@ def data_kind( # noqa: PLR0911
268274
):
269275
return "file"
270276

277+
# A StringIO object.
278+
if isinstance(data, io.StringIO):
279+
return "stringio"
280+
271281
# An option argument, mainly for dealing with optional virtual files
272282
if isinstance(data, bool | int | float) or (data is None and not required):
273283
return "arg"

0 commit comments

Comments
 (0)