Wrap legend

pygmt/base_plotting.py

@@ -11,7 +11,9 @@
     fmt_docstring,
     use_alias,
     kwargs_to_strings,
+    GMTTempFile,
 )
+import re
 
 
 class BasePlotting:
@@ -345,6 +347,8 @@ def plot(self, x=None, y=None, data=None, sizes=None, direction=None, **kwargs):
                 arg_str = " ".join([fname, build_arg_string(kwargs)])
                 lib.call_module("plot", arg_str)
 
+        return kwargs
+
     @fmt_docstring
     @use_alias(
         R="region",
@@ -530,3 +534,71 @@ def image(self, imagefile, **kwargs):
         with Session() as lib:
             arg_str = " ".join([imagefile, build_arg_string(kwargs)])
             lib.call_module("image", arg_str)
+
+    @fmt_docstring
+    @use_alias(R="region", J="projection")
+    @kwargs_to_strings(R="sequence")
+    def legend(self, spec, **kwargs):
+        """
+        Plot legends on maps. [NEED TO UPDATE FOR HANDLES/LABELS FUNCTIONALITY]
+
+        Makes legends that can be overlaid on maps. Reads specific legend-related information from an input file. Unless otherwise noted, annotations will be made using the primary annotation font and size in effect (i.e., FONT_ANNOT_PRIMARY).
+
+        Full option list at :gmt-docs:`legend.html`
+
+        {aliases}
+
+        Parameters
+        ----------
+        {J}
+        {R}
+        D: str
+            ``'[g|j|J|n|x]refpoint+wwidth[/height][+jjustify][+lspacing][+odx[/dy]]'``
+            Defines the reference point on the map for the legend.
+        F : bool or str
+            ``'[+cclearances][+gfill][+i[[gap/]pen]][+p[pen]][+r[radius]][+s[[dx/dy/][shade]]]'``
+            Without further options, draws a rectangular border around the
+            legend using **MAP_FRAME_PEN**.
+        """
+        kwargs = self._preprocess(**kwargs)
+        with Session() as lib:
+            kind = data_kind(spec, None, None, None)
+            with GMTTempFile() as tmp:
+                if kind == "file":
+                    specfile = spec
+                elif kind == "matrix":
+                    specfile = tmp.name
+                    handles = spec[0]
+                    labels = spec[1]
+                    with open(specfile, "w") as file:
+                        for h, text in zip(handles, labels):
+
+                            # SYMBOL and SIZE
+                            symbol, size = re.findall("\d*\D+", h["S"])
+
+                            # FILL
+                            if "G" not in h.keys():
+                                fill = "-"
+                            else:
+                                fill = h["G"]
+
+                            # PEN
+                            if "W" not in h.keys():
+                                pen = "-"
+                            elif not h["W"]:
+                                pen = "default,black"
+                            else:
+                                pen = h["W"]
+
+                            # Write a line to the specfile
+                            file.write(
+                                "S - {} {} {} {} - {}\n".format(
+                                    symbol, size, fill, pen, text
+                                )
+                            )
+                else:
+                    raise GMTInvalidInput(
+                        "Unrecognized data type: {}".format(type(spec))
+                    )
+                arg_str = " ".join([specfile, build_arg_string(kwargs)])
+                lib.call_module("legend", arg_str)

pygmt/clib/session.py

@@ -112,7 +112,7 @@ class Session:
     """
 
     # The minimum version of GMT required
-    required_version = "6.0.0rc4"
+    required_version = "6.0.0rc2"
 
     @property
     def session_pointer(self):

pygmt/tests/test_legend.py

@@ -0,0 +1,36 @@
+"""
+Tests for legend
+"""
+import pytest
+
+from .. import Figure
+
+
+@pytest.mark.mpl_image_compare
+def test_legend():
+    """
+    Create 3-entry legend.
+    """
+    fig = Figure()
+
+    h1 = fig.plot(
+        x=[-5],
+        y=[5],
+        region=[-10, 10, -5, 10],
+        projection="X3i/0",
+        frame="a",
+        style="a15p",
+        pen="1.5p,purple",
+    )
+
+    h2 = fig.plot(x=[0], y=[5], style="t10p", color="cyan")
+
+    h3 = fig.plot(x=[5], y=[5], style="d5p", color="yellow", pen=True)
+
+    fig.legend(
+        [[h1, h2, h3], ["I am a star!", "I am a triangle!", "I am a diamond!"]],
+        F=True,
+        D="g0/0+w2i+jCM",
+    )
+
+    return fig