Skip to content

Commit c2b3ee1

Browse files
committed
Figure.legend: Support passing a StringIO object as the legend specification
1 parent 486fce7 commit c2b3ee1

File tree

2 files changed

+19
-3
lines changed

2 files changed

+19
-3
lines changed

pygmt/src/legend.py

+6-3
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
legend - Plot a legend.
33
"""
44

5+
import io
56
import pathlib
67

78
from pygmt.clib import Session
@@ -30,7 +31,7 @@
3031
@kwargs_to_strings(R="sequence", c="sequence_comma", p="sequence")
3132
def legend(
3233
self,
33-
spec: str | pathlib.PurePath | None = None,
34+
spec: str | pathlib.PurePath | io.StringIO | None = None,
3435
position="JTR+jTR+o0.2c",
3536
box="+gwhite+p1p",
3637
**kwargs,
@@ -57,6 +58,7 @@ def legend(
5758
file
5859
- A string or a :class:`pathlib.PurePath` object pointing to the legend
5960
specification file
61+
- A :class:`io.StringIO` object containing the legend specification.
6062
6163
See :gmt-docs:`legend.html` for the definition of the legend specification.
6264
{projection}
@@ -89,10 +91,11 @@ def legend(
8991
kwargs["F"] = box
9092

9193
kind = data_kind(spec)
92-
if kind not in {"vectors", "file"}: # kind="vectors" means spec is None
94+
if kind not in {"vectors", "file", "stringio"}: # kind="vectors" means spec is None
9395
raise GMTInvalidInput(f"Unrecognized data type: {type(spec)}")
9496
if kind == "file" and is_nonstr_iter(spec):
9597
raise GMTInvalidInput("Only one legend specification file is allowed.")
9698

9799
with Session() as lib:
98-
lib.call_module(module="legend", args=build_arg_list(kwargs, infile=spec))
100+
with lib.virtualfile_in(data=spec, required_data=False) as vintbl:
101+
lib.call_module(module="legend", args=build_arg_list(kwargs, infile=vintbl))

pygmt/tests/test_legend.py

+13
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
Test Figure.legend.
33
"""
44

5+
import io
56
from pathlib import Path
67

78
import pytest
@@ -100,6 +101,18 @@ def test_legend_specfile(legend_spec):
100101
fig = Figure()
101102
fig.basemap(projection="x6i", region=[0, 1, 0, 1], frame=True)
102103
fig.legend(specfile.name, position="JTM+jCM+w5i")
104+
return fig
105+
106+
107+
@pytest.mark.mpl_image_compare(filename="test_legend_specfile.png")
108+
def test_legend_stringio(legend_spec):
109+
"""
110+
Test passing an legend specification via an io.StringIO object.
111+
"""
112+
spec = io.StringIO(legend_spec)
113+
fig = Figure()
114+
fig.basemap(projection="x6i", region=[0, 1, 0, 1], frame=True)
115+
fig.legend(spec, position="JTM+jCM+w5i")
103116
return fig
104117

105118

0 commit comments

Comments
 (0)