Handle figure legends as well as axes legends

Author: bcbnz

pgfutils.py

@@ -901,52 +901,59 @@ def save(figure: "matplotlib.figure.Figure|None" = None):
     # Figures to save.
     to_save = [("main_figure", figure, None)]
 
-    # Go through and fix up a few little quirks on the axes within this figure.
+    # Find all legends in this figure.
+    legends = list(figure.legends)
     for axes in figure.get_axes():
-        # Check if these axes have a legend.
         legend = axes.get_legend()
         if legend:
-            # Want to save the legend as a separate figure.
-            if config.pgfutils["separate_legend"]:
-                # Create a new figure to hold the legend with empty axes.
-                legend_fig = plt.figure()
-                legend_ax = legend_fig.add_subplot()
-                legend_ax.axis("off")
-
-                # Regenerate the legend on the new axes, allowing it to use the whole
-                # figure. Removing it from the original axes and then using add_artist
-                # on these axes seems to get the bounding box wrong.
-                legend_standalone = legend_ax.legend(
-                    *legend.axes.get_legend_handles_labels(),
-                    bbox_to_anchor=(0, 0, 1, 1),
-                    bbox_transform=legend_fig.transFigure,
-                    ncols=legend._ncols,
-                    numpoints=legend.numpoints,
-                    scatterpoints=legend.scatterpoints,
-                )
+            legends.append(legend)
+
+    # Configure the legends.
+    for legend in legends:
+        # Want to save the legend as a separate figure.
+        if config.pgfutils["separate_legend"]:
+            # Create a new figure to hold the legend with empty axes.
+            legend_fig = plt.figure()
+            legend_ax = legend_fig.add_subplot()
+            legend_ax.axis("off")
+
+            # Regenerate the legend on the new axes, allowing it to use the whole
+            # figure. Removing it from the original axes and then using add_artist
+            # on these axes seems to get the bounding box wrong.
+            legend_standalone = legend_ax.legend(
+                legend.legend_handles,
+                [text.get_text() for text in legend.texts],
+                bbox_to_anchor=(0, 0, 1, 1),
+                bbox_transform=legend_fig.transFigure,
+                ncols=legend._ncols,
+                numpoints=legend.numpoints,
+                scatterpoints=legend.scatterpoints,
+            )
 
-                # Measure its bounding box.
-                bbox = legend_standalone.get_tightbbox()
-                if not bbox:
-                    raise RuntimeError("could not determine legend bounding box")
+            # Measure its bounding box.
+            bbox = legend_standalone.get_tightbbox()
+            if not bbox:
+                raise RuntimeError("could not determine legend bounding box")
 
-                # This is in pixels; convert to inches as savefig() will need.
-                bbox = bbox.transformed(legend_fig.dpi_scale_trans.inverted())
+            # This is in pixels; convert to inches as savefig() will need.
+            bbox = bbox.transformed(legend_fig.dpi_scale_trans.inverted())
 
-                # Remove the original legend and replace the reference.
-                legend.remove()
-                legend = legend_standalone
+            # Remove the original legend and replace the reference.
+            legend.remove()
+            legend = legend_standalone
 
-                # And save the standalone figure as a legend.
-                to_save.append(("legend", legend_fig, bbox))
+            # And save the standalone figure as a legend.
+            to_save.append(("legend", legend_fig, bbox))
 
-            # There are no rcParams for the legend properties; set them directly.
-            frame = legend.get_frame()
-            frame.set_linewidth(config.pgfutils["legend_border_width"])
-            frame.set_alpha(config.pgfutils["legend_opacity"])
-            frame.set_ec(config.pgfutils["legend_border_color"])
-            frame.set_fc(config.pgfutils["legend_background"])
+        # There are no rcParams for the legend properties; set them directly.
+        frame = legend.get_frame()
+        frame.set_linewidth(config.pgfutils["legend_border_width"])
+        frame.set_alpha(config.pgfutils["legend_opacity"])
+        frame.set_ec(config.pgfutils["legend_border_color"])
+        frame.set_fc(config.pgfutils["legend_background"])
 
+    # Go through and fix up a few little quirks on the axes within this figure.
+    for axes in figure.get_axes():
         # Some PDF viewers show white lines through vector colorbars. This is a bug in
         # the viewers, but can be worked around by forcing the edge of the patches in
         # the colorbar to have the same colour as their faces.  This doesn't work with

tests/sources/legend/plot_with_separate_figure_legend.py

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: Blair Bonnett
+# SPDX-License-Identifier: BSD-3-Clause
+
+from pgfutils import save, setup_figure
+
+setup_figure(width=1, height=1, separate_legend=True)
+
+from matplotlib import pyplot as plt
+import numpy as np
+
+t = np.arange(0, 1, 0.005)
+s1 = np.sin(2 * np.pi * t)
+s2 = np.sin(4 * np.pi * t)
+s3 = np.sin(6 * np.pi * t)
+
+# Disable all other frames and sources of text.
+fig = plt.figure(frameon=False)
+ax = fig.add_axes([0, 0, 1, 1])
+ax.axis("off")
+
+(a1,) = ax.plot(t, s1)
+(a2,) = ax.plot(t, s2)
+(a3,) = ax.plot(t, s3)
+fig.legend([a1, a2, a3], ["One", "Two", "Three"])
+
+save()

tests/test_legend.py

@@ -151,3 +151,33 @@ def test_separate_legend():
         assert legend["stroke"] == approx([1, 0.5, 0]), "incorrect border colour"
         assert legend["strokeopacity"] == approx(0.7), "incorrect border opacity"
         assert legend["linewidth"] == approx(4, abs=0.1), "incorrect border width"
+
+
+def test_separate_figure_legend():
+    """Check plot with separate figure legend..."""
+    with build_pypgf(srcdir, "plot_with_separate_figure_legend.py") as res:
+        assert res.returncode == 0, (
+            f"Building {srcdir / 'plot_with_separate_figure_legend.pypgf'} failed."
+        )
+
+        # Should be no legend in main figure.
+        legend = extract_legend(srcdir / "plot_with_separate_figure_legend.pypgf")
+        assert legend is None, "legend retained in main figure"
+
+        # Get the legend from the separate output.
+        legend = extract_legend(
+            srcdir / "plot_with_separate_figure_legend_legend0.pypgf"
+        )
+
+        # Check the text sizes are correct.
+        assert legend["text_sizes"] == approx([14, 14, 14]), "incorrect font sizes"
+
+        # Now check the values are correct. Note the increased margin for the
+        # line size -- the output value often seems to be a wee way off the
+        # number. I'd guess this is due to some rounding in the exporter. At
+        # the end of the day 0.1 of a point is not that noticeable!
+        assert legend["fill"] == approx([0, 0.5, 1]), "incorrect background colour"
+        assert legend["fillopacity"] == approx(0.7), "incorrect background opacity"
+        assert legend["stroke"] == approx([1, 0.5, 0]), "incorrect border colour"
+        assert legend["strokeopacity"] == approx(0.7), "incorrect border opacity"
+        assert legend["linewidth"] == approx(4, abs=0.1), "incorrect border width"