Figure.legend: Support passing a StringIO object as the legend specification

Author: seisman

pygmt/src/legend.py

@@ -2,6 +2,9 @@
 legend - Plot a legend.
 """
 
+import io
+import pathlib
+
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -26,7 +29,13 @@
     t="transparency",
 )
 @kwargs_to_strings(R="sequence", c="sequence_comma", p="sequence")
-def legend(self, spec=None, position="JTR+jTR+o0.2c", box="+gwhite+p1p", **kwargs):
+def legend(
+    self,
+    spec: str | pathlib.PurePath | io.StringIO | None = None,
+    position="JTR+jTR+o0.2c",
+    box="+gwhite+p1p",
+    **kwargs,
+):
     r"""
     Plot legends on maps.
 
@@ -42,10 +51,15 @@ def legend(self, spec=None, position="JTR+jTR+o0.2c", box="+gwhite+p1p", **kwarg
 
     Parameters
     ----------
-    spec : None or str
-        Either ``None`` [Default] for using the automatically generated legend
-        specification file, or a *filename* pointing to the legend
-        specification file.
+    spec
+        The legend specification. It can be:
+
+        - ``None`` means using the automatically generated legend specification file.
+        - A string or a :class:`pathlib.PurePath` object pointing to the legend
+          specification file.
+        - A :class:`io.StringIO` object containing the legend specification.
+
+        See :gmt-docs:`legend.html` for the definition of the legend specification.
     {projection}
     {region}
     position : str
@@ -75,12 +89,12 @@ def legend(self, spec=None, position="JTR+jTR+o0.2c", box="+gwhite+p1p", **kwarg
         if kwargs.get("F") is None:
             kwargs["F"] = box
 
+    kind = data_kind(spec)
+    if kind not in {"vectors", "file", "stringio"}:  # kind="vectors" means spec is None
+        raise GMTInvalidInput(f"Unrecognized data type: {type(spec)}")
+    if kind == "file" and is_nonstr_iter(spec):
+        raise GMTInvalidInput("Only one legend specification file is allowed.")
+
     with Session() as lib:
-        if spec is None:
-            specfile = ""
-        elif data_kind(spec) == "file" and not is_nonstr_iter(spec):
-            # Is a file but not a list of files
-            specfile = spec
-        else:
-            raise GMTInvalidInput(f"Unrecognized data type: {type(spec)}")
-        lib.call_module(module="legend", args=build_arg_list(kwargs, infile=specfile))
+        with lib.virtualfile_in(data=spec, required_data=False) as vintbl:
+            lib.call_module(module="legend", args=build_arg_list(kwargs, infile=vintbl))

pygmt/tests/test_legend.py

@@ -2,6 +2,7 @@
 Test Figure.legend.
 """
 
+import io
 from pathlib import Path
 
 import pytest
@@ -10,12 +11,44 @@
 from pygmt.helpers import GMTTempFile
 
 
+@pytest.fixture(scope="module", name="legend_spec")
+def fixture_legend_spec():
+    """
+    A string contains a legend specification.
+    """
+    return """
+G -0.1i
+H 24 Times-Roman My Map Legend
+D 0.2i 1p
+N 2
+V 0 1p
+S 0.1i c 0.15i p300/12 0.25p 0.3i This circle is hachured
+S 0.1i e 0.15i yellow 0.25p 0.3i This ellipse is yellow
+S 0.1i w 0.15i green 0.25p 0.3i This wedge is green
+S 0.1i f0.1i+l+t 0.25i blue 0.25p 0.3i This is a fault
+S 0.1i - 0.15i - 0.25p,- 0.3i A dashed contour
+S 0.1i v0.1i+a40+e 0.25i magenta 0.25p 0.3i This is a vector
+S 0.1i i 0.15i cyan 0.25p 0.3i This triangle is boring
+V 0 1p
+D 0.2i 1p
+N 1
+G 0.05i
+G 0.05i
+G 0.05i
+L 9 4 R Smith et al., @%5%J. Geophys. Res., 99@%%, 2000
+G 0.1i
+P
+T Let us just try some simple text that can go on a few lines.
+T There is no easy way to predetermine how many lines will be required,
+T so we may have to adjust the box height to get the right size box.
+"""
+
+
 @pytest.mark.mpl_image_compare
 def test_legend_position():
     """
-    Test that plots a position with each of the four legend coordinate systems.
+    Test positioning the legend with different coordinate systems.
     """
-
     fig = Figure()
     fig.basemap(region=[-2, 2, -2, 2], frame=True)
     positions = ["jTR+jTR", "g0/1", "n0.2/0.2", "x4i/2i/2i"]
@@ -30,22 +63,18 @@ def test_legend_default_position():
     """
     Test using the default legend position.
     """
-
     fig = Figure()
-
     fig.basemap(region=[-1, 1, -1, 1], frame=True)
-
     fig.plot(x=[0], y=[0], style="p10p", label="Default")
     fig.legend()
-
     return fig
 
 
 @pytest.mark.benchmark
 @pytest.mark.mpl_image_compare
 def test_legend_entries():
     """
-    Test different marker types/shapes.
+    Test legend using the automatically generated legend entries.
     """
     fig = Figure()
     fig.basemap(projection="x1i", region=[0, 7, 3, 7], frame=True)
@@ -59,48 +88,31 @@ def test_legend_entries():
     fig.plot(data="@Table_5_11.txt", pen="1.5p,gray", label="My lines")
     fig.plot(data="@Table_5_11.txt", style="t0.15i", fill="orange", label="Oranges")
     fig.legend(position="JTR+jTR")
-
     return fig
 
 
 @pytest.mark.mpl_image_compare
-def test_legend_specfile():
+def test_legend_specfile(legend_spec):
     """
-    Test specfile functionality.
+    Test passing a legend specification file.
     """
-
-    specfile_contents = """
-G -0.1i
-H 24 Times-Roman My Map Legend
-D 0.2i 1p
-N 2
-V 0 1p
-S 0.1i c 0.15i p300/12 0.25p 0.3i This circle is hachured
-S 0.1i e 0.15i yellow 0.25p 0.3i This ellipse is yellow
-S 0.1i w 0.15i green 0.25p 0.3i This wedge is green
-S 0.1i f0.1i+l+t 0.25i blue 0.25p 0.3i This is a fault
-S 0.1i - 0.15i - 0.25p,- 0.3i A dashed contour
-S 0.1i v0.1i+a40+e 0.25i magenta 0.25p 0.3i This is a vector
-S 0.1i i 0.15i cyan 0.25p 0.3i This triangle is boring
-V 0 1p
-D 0.2i 1p
-N 1
-G 0.05i
-G 0.05i
-G 0.05i
-L 9 4 R Smith et al., @%5%J. Geophys. Res., 99@%%, 2000
-G 0.1i
-P
-T Let us just try some simple text that can go on a few lines.
-T There is no easy way to predetermine how many lines will be required,
-T so we may have to adjust the box height to get the right size box.
-"""
-
     with GMTTempFile() as specfile:
-        Path(specfile.name).write_text(specfile_contents, encoding="utf-8")
+        Path(specfile.name).write_text(legend_spec, encoding="utf-8")
         fig = Figure()
         fig.basemap(projection="x6i", region=[0, 1, 0, 1], frame=True)
         fig.legend(specfile.name, position="JTM+jCM+w5i")
+        return fig
+
+
+@pytest.mark.mpl_image_compare(filename="test_legend_specfile.png")
+def test_legend_stringio(legend_spec):
+    """
+    Test passing an legend specification via an io.StringIO object.
+    """
+    spec = io.StringIO(legend_spec)
+    fig = Figure()
+    fig.basemap(projection="x6i", region=[0, 1, 0, 1], frame=True)
+    fig.legend(spec, position="JTM+jCM+w5i")
     return fig
 
 
@@ -111,3 +123,6 @@ def test_legend_fails():
     fig = Figure()
     with pytest.raises(GMTInvalidInput):
         fig.legend(spec=["@Table_5_11.txt"])
+
+    with pytest.raises(GMTInvalidInput):
+        fig.legend(spec=[1, 2])