ENH: add outside kwarg to figure legend

jklymak · jklymak · commit c6a986994388 · 2021-03-21T09:38:40.000-07:00
diff --git a/doc/users/next_whats_new/legend-figure-outside.rst b/doc/users/next_whats_new/legend-figure-outside.rst
@@ -0,0 +1,8 @@
+:orphan:
+
+Figure legends now have *outside* keyword argument
+--------------------------------------------------
+If a legend is made on a figure (or subfigure), and constrained_layout is being used, 
+then setting the *outside* kwarg to *True* on `.figure.legend` will move axes to 
+make room for the legend.  See :doc:`/tutorials/intermediate/legend_guide` for an 
+example.  
diff --git a/examples/text_labels_and_annotations/figlegend_demo.py b/examples/text_labels_and_annotations/figlegend_demo.py
@@ -28,3 +28,26 @@
 
 plt.tight_layout()
 plt.show()
+
+##############################################################################
+# Sometimes we do not want the legend to overlap the axes.  If you use
+# constrained_layout you can specify ``outside=True``, the legend will
+# not overlap.
+
+fig, axs = plt.subplots(1, 2, constrained_layout=True)
+
+x = np.arange(0.0, 2.0, 0.02)
+y1 = np.sin(2 * np.pi * x)
+y2 = np.exp(-x)
+l1, = axs[0].plot(x, y1)
+l2, = axs[0].plot(x, y2, marker='o')
+
+y3 = np.sin(4 * np.pi * x)
+y4 = np.exp(-2 * x)
+l3, = axs[1].plot(x, y3, color='tab:green')
+l4, = axs[1].plot(x, y4, color='tab:red', marker='^')
+
+fig.legend((l1, l2), ('Line 1', 'Line 2'), 'upper left')
+fig.legend((l3, l4), ('Line 3', 'Line 4'), 'upper right', outside=True)
+
+plt.show()
diff --git a/lib/matplotlib/_constrained_layout.py b/lib/matplotlib/_constrained_layout.py
@@ -273,6 +273,27 @@ def _make_layout_margins(fig, renderer, *, w_pad=0, h_pad=0,
         # pass the new margins down to the layout grid for the solution...
         gs._layoutgrid.edit_outer_margin_mins(margin, ss)
 
+    # make margins for figure-level legends:
+    for leg in fig.legends:
+        inv_trans_fig = None
+        if leg._outside and leg._bbox_to_anchor is None:
+            if inv_trans_fig is None:
+                inv_trans_fig = fig.transFigure.inverted().transform_bbox
+            bbox = inv_trans_fig(leg.get_tightbbox(renderer))
+            w = bbox.width + 2 * w_pad
+            h = bbox.height + 2 * h_pad
+            margin = 'right'
+            if ((leg._loc in (3, 4) and leg._outside == 'lower') or
+                    (leg._loc == 8)):
+                fig._layoutgrid.edit_margin_min('bottom', h)
+            elif ((leg._loc in (1, 2) and leg._outside == 'upper') or
+                    (leg._loc == 9)):
+                fig._layoutgrid.edit_margin_min('top', h)
+            elif leg._loc in (1, 4, 5, 7):
+                fig._layoutgrid.edit_margin_min('right', w)
+            elif leg._loc in (2, 3, 6):
+                fig._layoutgrid.edit_margin_min('left', w)
+
 
 def _make_margin_suptitles(fig, renderer, *, w_pad=0, h_pad=0):
     # Figure out how large the suptitle is and make the
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -1085,6 +1085,7 @@ def legend(self, *args, **kwargs):
         # explicitly set the bbox transform if the user hasn't.
         l = mlegend.Legend(self, handles, labels, *extra_args,
                            bbox_transform=transform, **kwargs)
+        l._outside = kwargs.pop('outside', False)
         self.legends.append(l)
         l._remove_method = self.legends.remove
         self.stale = True
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -263,6 +263,13 @@ def _update_bbox_to_anchor(self, loc_in_canvas):
     The custom dictionary mapping instances or types to a legend
     handler. This *handler_map* updates the default handler map
     found at `matplotlib.legend.Legend.get_legend_handler_map`.
+
+outside : bool or string
+    For `.figure.legend` when used with `.figure.set_constrained_layout`.
+    If True, place the legend outside all the axes in the figure.
+    loc='upper right/left' wil usually put beside the figure, but if
+    outside='upper', then place above the axes; if loc='lower right/left'
+    and outside='lower' then place below the axes.
 """)
 
 
@@ -332,6 +339,7 @@ def __init__(self, parent, handles, labels,
                  bbox_transform=None,  # transform for the bbox
                  frameon=None,  # draw frame
                  handler_map=None,
+                 outside=False,
                  ):
         """
         Parameters
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -3,6 +3,7 @@
 from unittest import mock
 
 import numpy as np
+from numpy.testing import assert_allclose
 import pytest
 
 from matplotlib.testing.decorators import image_comparison
@@ -380,6 +381,51 @@ def test_warn_args_kwargs(self):
             "be discarded.")
 
 
+def test_figure_legend_outside():
+    outside = [True]*9 + ['upper', 'upper', 'lower', 'lower']
+    print(outside)
+    todos = [1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4]
+    axbb = [[20.347556,  27.722556, 659.249, 588.833],     # upper right
+            [151.681556, 27.722556, 790.583, 588.833],     # upper left
+            [151.681556, 27.722556, 790.583, 588.833],     # lower left
+            [20.347556,  27.722556, 659.249, 588.833],     # lower right
+            [20.347556,  27.722556, 659.249, 588.833],     # right
+            [151.681556, 27.722556, 790.583, 588.833],     # center left
+            [20.347556,  27.722556, 659.249, 588.833],     # center right
+            [20.347556,  71.056556, 790.583, 588.833],     # lower center
+            [20.347556,  27.722556, 790.583, 545.499],  # upper center
+            [20.347556,  27.722556, 790.583, 545.499],  # up-right,'upper'
+            [20.347556,  27.722556, 790.583, 545.499],  # up-left,'upper'
+            [20.347556,  71.056556, 790.583, 588.833],  # low-left,'lower'
+            [20.347556,  71.056556, 790.583, 588.833],  # low-right,'lower'
+            ]
+    legbb = [[667., 555., 790., 590.],  # upper right
+             [10., 555., 133., 590.],   # upper left
+             [10.,  10., 133.,  45.],   # lower left
+             [667,  10., 790.,  45.],   # lower right
+             [667., 282.5, 790., 317.5],
+             [10., 282.5, 133., 317.5],
+             [667., 282.5, 790., 317.5],
+             [338.5,  10., 461.5,  45.],
+             [338.5, 555., 461.5, 590.],
+             [667., 555., 790., 590.],  # upper right
+             [10., 555., 133., 590.],   # upper left
+             [10.,  10., 133.,  45.],   # lower left
+             [667,  10., 790.,  45.],   # lower right
+             ]
+    for nn, todo in enumerate(todos):
+        print(todo)
+        fig, axs = plt.subplots(constrained_layout=True, dpi=100)
+        axs.plot(range(10), label='Boo1')
+        leg = fig.legend(loc=todo, outside=outside[nn])
+        renderer = fig.canvas.get_renderer()
+        fig.canvas.draw()
+        assert_allclose(axs.get_window_extent(renderer=renderer).extents,
+                        axbb[nn])
+        assert_allclose(leg.get_window_extent(renderer=renderer).extents,
+                        legbb[nn])
+
+
 @image_comparison(['legend_stackplot.png'])
 def test_legend_stackplot():
     """Test legend for PolyCollection using stackplot."""
diff --git a/tutorials/intermediate/legend_guide.py b/tutorials/intermediate/legend_guide.py
@@ -115,20 +115,43 @@
 #
 # More examples of custom legend placement:
 
-plt.subplot(211)
-plt.plot([1, 2, 3], label="test1")
-plt.plot([3, 2, 1], label="test2")
+fig, axs = plt.subplot_mosaic([['top', 'top'], ['left', 'right']])
+
+axs['right'].remove()
+
+axs['top'].plot([1, 2, 3], label="test1")
+axs['top'].plot([3, 2, 1], label="test2")
 
 # Place a legend above this subplot, expanding itself to
 # fully use the given bounding box.
-plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc='lower left',
-           ncol=2, mode="expand", borderaxespad=0.)
+axs['top'].legend(bbox_to_anchor=(0., 1.02, 1., .102), loc='lower left',
+                  ncol=2, mode="expand", borderaxespad=0.)
+
+axs['left'].plot([1, 2, 3], label="test1")
+axs['left'].plot([3, 2, 1], label="test2")
+# Place a legend to the right of this smaller subplot.
+axs['left'].legend(bbox_to_anchor=(1.05, 1), loc='upper left',
+                   borderaxespad=0.)
+
+plt.show()
+
+##############################################################################
+# Figure legends
+# --------------
+#
+# Sometimes it makes more sense to place the axes relative to the (sub)figure
+# rather than individual axes.  By using ``constrained_layout`` and
+# ``outside=True`` the legend is drawn outside the axes on the (sub)figure.
+
+fig, axs = plt.subplot_mosaic([['left', 'right']], constrained_layout=True)
+
+axs['left'].plot([1, 2, 3], label="test1")
+axs['left'].plot([3, 2, 1], label="test2")
 
-plt.subplot(223)
-plt.plot([1, 2, 3], label="test1")
-plt.plot([3, 2, 1], label="test2")
+axs['right'].plot([1, 2, 3], 'C2', label="test3")
+axs['right'].plot([3, 2, 1], 'C3', label="test4")
 # Place a legend to the right of this smaller subplot.
-plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0.)
+fig.legend(loc='upper right', outside=True)
 
 plt.show()
 
