Merge pull request #12455 from kbrose/warn-on-legend-slowness

NelleV · web-flow · commit 5702999ffd64 · 2019-01-20T22:53:03.000-08:00
Warn when "best" loc of legend is used with lots of data
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -484,6 +484,7 @@ def __init__(self, parent, handles, labels,
             raise TypeError("Legend needs either Axes or Figure as parent")
         self.parent = parent
 
+        self._loc_used_default = loc is None
         if loc is None:
             loc = rcParams["legend.loc"]
             if not self.isaxes and loc in [0, 'best']:
@@ -568,7 +569,10 @@ def __init__(self, parent, handles, labels,
         else:
             self.get_frame().set_alpha(framealpha)
 
-        self._loc = loc
+        tmp = self._loc_used_default
+        self._set_loc(loc)
+        self._loc_used_default = tmp  # ignore changes done by _set_loc
+
         # figure out title fontsize:
         if title_fontsize is None:
             title_fontsize = rcParams['legend.title_fontsize']
@@ -592,6 +596,7 @@ def _set_loc(self, loc):
         # find_offset function will be provided to _legend_box and
         # _legend_box will draw itself at the location of the return
         # value of the find_offset.
+        self._loc_used_default = False
         self._loc_real = loc
         self.stale = True
         self._legend_box.set_offset(self._findoffset)
@@ -1108,6 +1113,12 @@ def _find_best_position(self, width, height, renderer, consider=None):
         assert self.isaxes
 
         verts, bboxes, lines, offsets = self._auto_legend_data()
+        if self._loc_used_default and verts.shape[0] > 200000:
+            # this size results in a 3+ second render time on a good machine
+            cbook._warn_external(
+                'Creating legend with loc="best" can be slow with large '
+                'amounts of data.'
+            )
 
         bbox = Bbox.from_bounds(0, 0, width, height)
         if consider is None:
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -13,6 +13,7 @@
 from matplotlib.legend_handler import HandlerTuple
 import matplotlib.legend as mlegend
 from matplotlib.cbook.deprecation import MatplotlibDeprecationWarning
+from matplotlib import rc_context
 
 
 def test_legend_ordereddict():
@@ -545,3 +546,28 @@ def test_alpha_handles():
         lh.set_alpha(1.0)
     assert lh.get_facecolor()[:-1] == hh[1].get_facecolor()[:-1]
     assert lh.get_edgecolor()[:-1] == hh[1].get_edgecolor()[:-1]
+
+
+def test_warn_big_data_best_loc():
+    fig, ax = plt.subplots()
+    ax.plot(np.arange(200001), label='Is this big data?')
+    with pytest.warns(UserWarning) as records:
+        with rc_context({'legend.loc': 'best'}):
+            l = ax.legend()
+        fig.canvas.draw()
+    # 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():
+    fig, ax = plt.subplots()
+    ax.plot(np.arange(200001), label='Is this big data?')
+    with pytest.warns(None) as records:
+        l = ax.legend('best')
+        fig.canvas.draw()
+    assert len(records) == 0
