Update stem linecolor in legend

dstansby · dstansby · commit 16d88d8494c3 · 2018-06-28T10:29:19.000+01:00
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -589,7 +589,6 @@ def get_ydata(self, legend, xdescent, ydescent, width, height, fontsize):
     def create_artists(self, legend, orig_handle,
                        xdescent, ydescent, width, height, fontsize,
                        trans):
-
         markerline, stemlines, baseline = orig_handle
 
         xdata, xdata_marker = self.get_xdata(legend, xdescent, ydescent,
@@ -607,9 +606,18 @@ def create_artists(self, legend, orig_handle,
         self.update_prop(leg_markerline, markerline, legend)
 
         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
 
         leg_baseline = Line2D([np.min(xdata), np.max(xdata)],
                               [bottom, bottom])
@@ -625,6 +633,13 @@ def create_artists(self, legend, orig_handle,
 
         return artists
 
+    def _copy_collection_props(self, legend_handle, orig_handle):
+        '''
+        Method to copy properties from a LineCollection (orig_handle) to a
+        Line2D (legend_handle).
+        '''
+        legend_handle._color = orig_handle.get_color()[0]
+
 
 class HandlerTuple(HandlerBase):
     """
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/stem.png b/lib/matplotlib/tests/baseline_images/test_axes/stem.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -2928,7 +2928,7 @@ def test_stem():
     x = np.linspace(0.1, 2 * np.pi, 100)
 
     fig, ax = plt.subplots()
-    ax.stem(x, np.cos(x), linefmt='C0-', markerfmt='k+', basefmt='C1-.',
+    ax.stem(x, np.cos(x), linefmt='C2-', markerfmt='k+', basefmt='C1-.',
             label='Stem')
     ax.legend()
 
