Merge pull request #26483 from QuLogic/loc-best-warning

ksunden · web-flow · commit fa3aea9f262a · 2023-08-10T12:40:38.000-05:00
Improve legend(loc='best') warning and test
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -1158,30 +1158,22 @@ def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer):
             loc, bbox, parentbbox,
             self.borderaxespad * renderer.points_to_pixels(self._fontsize))
 
-    def _find_best_position(self, width, height, renderer, consider=None):
-        """
-        Determine the best location to place the legend.
-
-        *consider* is a list of ``(x, y)`` pairs to consider as a potential
-        lower-left corner of the legend. All are display coords.
-        """
+    def _find_best_position(self, width, height, renderer):
+        """Determine the best location to place the legend."""
         assert self.isaxes  # always holds, as this is only called internally
 
         start_time = time.perf_counter()
 
         bboxes, lines, offsets = self._auto_legend_data()
 
         bbox = Bbox.from_bounds(0, 0, width, height)
-        if consider is None:
-            consider = [self._get_anchored_bbox(x, bbox,
-                                                self.get_bbox_to_anchor(),
-                                                renderer)
-                        for x in range(1, len(self.codes))]
 
         candidates = []
-        for idx, (l, b) in enumerate(consider):
+        for idx in range(1, len(self.codes)):
+            l, b = self._get_anchored_bbox(idx, bbox,
+                                           self.get_bbox_to_anchor(),
+                                           renderer)
             legendBox = Bbox.from_bounds(l, b, width, height)
-            badness = 0
             # XXX TODO: If markers are present, it would be good to take them
             # into account when checking vertex overlaps in the next line.
             badness = (sum(legendBox.count_contains(line.vertices)
@@ -1190,10 +1182,10 @@ def _find_best_position(self, width, height, renderer, consider=None):
                        + legendBox.count_overlaps(bboxes)
                        + sum(line.intersects_bbox(legendBox, filled=False)
                              for line in lines))
-            if badness == 0:
-                return l, b
             # Include the index to favor lower codes in case of a tie.
             candidates.append((badness, idx, (l, b)))
+            if badness == 0:
+                break
 
         _, _, (l, b) = min(candidates)
 
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -1,5 +1,7 @@
 import collections
+import itertools
 import platform
+import time
 from unittest import mock
 import warnings
 
@@ -1109,29 +1111,43 @@ def test_usetex_no_warn(caplog):
     assert "Font family ['serif'] not found." not in caplog.text
 
 
-def test_warn_big_data_best_loc():
+def test_warn_big_data_best_loc(monkeypatch):
+    # Force _find_best_position to think it took a long time.
+    counter = itertools.count(0, step=1.5)
+    monkeypatch.setattr(time, 'perf_counter', lambda: next(counter))
+
     fig, ax = plt.subplots()
     fig.canvas.draw()  # So that we can call draw_artist later.
-    for idx in range(1000):
-        ax.plot(np.arange(5000), label=idx)
+
+    # Place line across all possible legend locations.
+    x = [0.9, 0.1, 0.1, 0.9, 0.9, 0.5]
+    y = [0.95, 0.95, 0.05, 0.05, 0.5, 0.5]
+    ax.plot(x, y, 'o-', label='line')
+
     with rc_context({'legend.loc': 'best'}):
         legend = ax.legend()
-    with pytest.warns(UserWarning) as records:
+    with pytest.warns(UserWarning,
+                      match='Creating legend with loc="best" can be slow with large '
+                      'amounts of data.') as records:
         fig.draw_artist(legend)  # Don't bother drawing the lines -- it's slow.
     # The _find_best_position method of Legend is called twice, duplicating
     # the warning message.
     assert len(records) == 2
-    for record in records:
-        assert str(record.message) == (
-            'Creating legend with loc="best" can be slow with large '
-            'amounts of data.')
 
 
-def test_no_warn_big_data_when_loc_specified():
+def test_no_warn_big_data_when_loc_specified(monkeypatch):
+    # Force _find_best_position to think it took a long time.
+    counter = itertools.count(0, step=1.5)
+    monkeypatch.setattr(time, 'perf_counter', lambda: next(counter))
+
     fig, ax = plt.subplots()
     fig.canvas.draw()
-    for idx in range(1000):
-        ax.plot(np.arange(5000), label=idx)
+
+    # Place line across all possible legend locations.
+    x = [0.9, 0.1, 0.1, 0.9, 0.9, 0.5]
+    y = [0.95, 0.95, 0.05, 0.05, 0.5, 0.5]
+    ax.plot(x, y, 'o-', label='line')
+
     legend = ax.legend('best')
     fig.draw_artist(legend)  # Check that no warning is emitted.
 
