Do deprecation better

dstansby · dstansby · commit ec8f5f7ed764 · 2018-10-04T10:03:46.000+01:00
diff --git a/doc/api/api_changes/2017-10-24-DS.rst b/doc/api/api_changes/2017-10-24-DS.rst
diff --git a/doc/api/next_api_changes/2017-10-24-DS.rst b/doc/api/next_api_changes/2017-10-24-DS.rst
@@ -0,0 +1,14 @@
+`StemContainer` now stores `LineCollection`
+-------------------------------------------
+
+`StemContainer` objects can now store a `LineCollection` object instead of a
+list of `Line2D` objects for stem lines plotted using `ax.stem`. This gives a
+very large performance boost to displaying and moving `ax.stem` plots.
+
+This will become the default behaviour in Matplotlib 3.3. To use it now, the
+``use_line_collection`` keyword argument to ~`.axes.stem` can be set to
+``True``.
+
+Individual line segments can be extracted from the `LineCollection` using
+`LineCollection.get_segements()`. See the `LineCollection` documentation for
+other methods to retrieve the collection properties.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2524,7 +2524,7 @@ def broken_barh(self, xranges, yrange, **kwargs):
 
     @_preprocess_data(replace_all_args=True, label_namer=None)
     def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None,
-             bottom=0, label=None):
+             bottom=0, label=None, use_line_collection=False):
         """
         Create a stem plot.
 
@@ -2583,6 +2583,12 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None,
         label : str, optional, default: None
             The label to use for the stems in legends.
 
+        use_line_collection : bool, optional, default: False
+            If ``True``, store and plot the stem lines as a
+            ~`.collections.LineCollection` instead of individual lines. This
+            significantly increases performance, and will become the default
+            option in Matplotlib 3.3. If ``False``, defaults to old behaviour.
+
 
         Returns
         -------
@@ -2665,24 +2671,40 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None,
         else:
             basestyle, basemarker, basecolor = _process_plot_format(basefmt)
 
+        # New behaviour in 3.1 is to use a LineCollection for the stemlines
+        if use_line_collection:
+            stemlines = []
+            for thisx, thisy in zip(x, y):
+                stemlines.append(((thisx, bottom), (thisx, thisy)))
+            stemlines = mcoll.LineCollection(stemlines, linestyles=linestyle,
+                                             colors=linecolor,
+                                             label='_nolegend_')
+            self.add_collection(stemlines)
+        # Old behaviour is to plot each of the lines individually
+        else:
+            warnings.warn(
+                'In Matplotlib 3.3 individual lines on a stem plot will be '
+                'added as a LineCollection instead of individual lines.\n'
+                'This significantly improves the performance of a stem plot.\n'
+                'To remove this warning and switch to the new behaviour, '
+                'set the "use_line_collection" keyword argument to True.')
+            stemlines = []
+            for thisx, thisy in zip(x, y):
+                l, = self.plot([thisx, thisx], [bottom, thisy],
+                               color=linecolor, linestyle=linestyle,
+                               marker=linemarker, label="_nolegend_")
+            stemlines.append(l)
+
         markerline, = self.plot(x, y, color=markercolor, linestyle=markerstyle,
                                 marker=markermarker, label="_nolegend_")
 
-        lines = []
-        for thisx, thisy in zip(x, y):
-            lines.append(((thisx, bottom), (thisx, thisy)))
-        stemlines = mcoll.LineCollection(lines, linestyles=linestyle,
-                                         colors=linecolor, label='_nolegend_')
-        self.add_collection(stemlines)
-
         baseline, = self.plot([np.min(x), np.max(x)], [bottom, bottom],
                               color=basecolor, linestyle=basestyle,
                               marker=basemarker, label="_nolegend_")
 
         stem_container = StemContainer((markerline, stemlines, baseline),
                                        label=label)
         self.add_container(stem_container)
-
         return stem_container
 
     @_preprocess_data(replace_names=["x", "explode", "labels", "colors"],
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -590,6 +590,10 @@ def create_artists(self, legend, orig_handle,
                        xdescent, ydescent, width, height, fontsize,
                        trans):
         markerline, stemlines, baseline = orig_handle
+        # Check to see if the stemcontainer is storing lines as a list or a
+        # LineCollection. Eventually using a list will be removed, and this
+        # logic can also be removed.
+        using_linecoll = isinstance(stemlines, mcoll.LineCollection)
 
         xdata, xdata_marker = self.get_xdata(legend, xdescent, ydescent,
                                              width, height, fontsize)
@@ -604,17 +608,24 @@ def create_artists(self, legend, orig_handle,
 
         leg_stemlines = []
 
-        # update_prop() usually takes two Line2D collections;
-        # override temporarily to copy properties from a LineCollection
-        orig_update_func = self._update_prop_func
-        self._update_prop_func = self._copy_collection_props
-
-        for thisx, thisy in zip(xdata_marker, ydata):
-            thisline = Line2D([thisx, thisx], [bottom, thisy])
-            leg_stemlines.append(thisline)
-            self.update_prop(thisline, stemlines, legend)
-        # Reset update_prop_func
-        self._update_prop_func = orig_update_func
+        if using_linecoll:
+            # update_prop() usually takes two Line2D collections;
+            # override temporarily to copy properties from a LineCollection
+            orig_update_func = self._update_prop_func
+            self._update_prop_func = self._copy_collection_props
+
+            for thisx, thisy in zip(xdata_marker, ydata):
+                thisline = Line2D([thisx, thisx], [bottom, thisy])
+                leg_stemlines.append(thisline)
+                self.update_prop(thisline, stemlines, legend)
+
+            self._update_prop_func = orig_update_func
+        else:
+            for thisx, thisy in zip(xdata_marker, ydata):
+                thisline = Line2D([thisx, thisx], [bottom, thisy])
+                leg_stemlines.append(thisline)
+            for lm, m in zip(leg_stemlines, stemlines):
+                self.update_prop(lm, m, legend)
 
         leg_baseline = Line2D([np.min(xdata), np.max(xdata)],
                               [bottom, bottom])
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -2958,14 +2958,23 @@ def test_hist_stacked_weighted():
     ax.hist((d1, d2), weights=(w1, w2), histtype="stepfilled", stacked=True)
 
 
+@pytest.mark.parametrize("use_line_collection", [True, False],
+                         ids=['w/ line collection', 'w/o line collection'])
 @image_comparison(baseline_images=['stem'], extensions=['png'], style='mpl20',
                   remove_text=True)
-def test_stem():
+def test_stem(use_line_collection):
     x = np.linspace(0.1, 2 * np.pi, 100)
+    args = (x, np.cos(x))
+    kwargs = dict(linefmt='C2-', markerfmt='k+', basefmt='C1-.',
+                  label='Stem', use_line_collection=use_line_collection)
 
     fig, ax = plt.subplots()
-    ax.stem(x, np.cos(x), linefmt='C2-', markerfmt='k+', basefmt='C1-.',
-            label='Stem')
+    if use_line_collection:
+        ax.stem(*args, **kwargs)
+    else:
+        with pytest.warns(UserWarning):
+            ax.stem(*args, **kwargs)
+
     ax.legend()
 
 
