Fix legend auto-placement in the presence of text or PolyCollections

alanlau28 · alextanned · rcomer · commit 87d05aea59a3 · 2024-01-21T13:03:59.000Z
Co-authored-by: alextanned &lt;74106760+alextanned@users.noreply.github.com&gt;
diff --git a/doc/api/next_api_changes/behavior/27469-AL.rst b/doc/api/next_api_changes/behavior/27469-AL.rst
@@ -0,0 +1,11 @@
+``loc='best'`` for ``legend`` now considers ``Text`` and ``PolyCollections``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The location selection ``legend`` now considers the existence of ``Text``
+and ``PolyCollections`` in the ``badness`` calculation.
+
+Note: The ``best`` option can already be quite slow for plots with large
+amounts of data. For ``PolyCollections``, it only considers the ``Path``
+of ``PolyCollections`` and not the enclosed area when checking for overlap
+to reduce additional latency. However, it can still be quite slow when
+there are large amounts of ``PolyCollections`` in the plot to check for.
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -978,10 +978,15 @@ def _auto_legend_data(self):
             elif isinstance(artist, Patch):
                 lines.append(
                     artist.get_transform().transform_path(artist.get_path()))
+            elif isinstance(artist, PolyCollection):
+                lines.extend(artist.get_transform().transform_path(path)
+                             for path in artist.get_paths())
             elif isinstance(artist, Collection):
                 transform, transOffset, hoffsets, _ = artist._prepare_points()
                 if len(hoffsets):
                     offsets.extend(transOffset.transform(hoffsets))
+            elif isinstance(artist, Text):
+                bboxes.append(artist.get_window_extent())
 
         return bboxes, lines, offsets
 
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/legend_stackplot.png b/lib/matplotlib/tests/baseline_images/test_legend/legend_stackplot.png
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -1391,3 +1391,39 @@ def test_legend_nolabels_draw():
     plt.plot([1, 2, 3])
     plt.legend()
     assert plt.gca().get_legend() is not None
+
+
+def test_legend_loc_polycollection():
+    # Test that the legend is placed in the correct
+    # position for 'best' for polycollection
+    x = [3, 4, 5]
+    y1 = [1, 1, 1]
+    y2 = [5, 5, 5]
+    leg_bboxes = []
+    fig, axs = plt.subplots(ncols=2, figsize=(10, 5))
+    for ax, loc in zip(axs.flat, ('best', 'lower left')):
+        ax.fill_between(x, y1, y2, color='gray', alpha=0.5, label='Shaded Area')
+        ax.set_xlim(0, 6)
+        ax.set_ylim(-1, 5)
+        leg = ax.legend(loc=loc)
+        fig.canvas.draw()
+        leg_bboxes.append(
+            leg.get_window_extent().transformed(ax.transAxes.inverted()))
+    assert_allclose(leg_bboxes[1].bounds, leg_bboxes[0].bounds)
+
+
+def test_legend_text():
+    # Test that legend is place in the correct
+    # position for 'best' when there is text in figure
+    fig, axs = plt.subplots(ncols=2, figsize=(10, 5))
+    leg_bboxes = []
+    for ax, loc in zip(axs.flat, ('best', 'lower left')):
+        x = [1, 2]
+        y = [2, 1]
+        ax.plot(x, y, label='plot name')
+        ax.text(1.5, 2, 'some text blahblah', verticalalignment='top')
+        leg = ax.legend(loc=loc)
+        fig.canvas.draw()
+        leg_bboxes.append(
+            leg.get_window_extent().transformed(ax.transAxes.inverted()))
+    assert_allclose(leg_bboxes[1].bounds, leg_bboxes[0].bounds)
