Merge pull request matplotlib#3303 from fcolas/master

ENH : Adding legend handler to PolyCollection and labels to stackplot

Author: tacaswell

doc/users/whats_new.rst

@@ -39,6 +39,16 @@ All selectors now implement ``set_active`` and ``get_active`` methods (also
 called when accessing the ``active`` property) to properly update and query
 whether they are active.
 
+
+New plotting features
+---------------------
+
+Support for legend for PolyCollection and stackplot
+```````````````````````````````````````````````````
+Added a `legend_handler` for :class:`~matplotlib.collections.PolyCollection` as well as a `labels` argument to
+:func:`~matplotlib.axes.Axes.stackplot`.
+
+
 .. _whats-new-1-4:
 
 new in matplotlib-1.4
@@ -280,6 +290,7 @@ polar-plot r-tick locations
 Added the ability to control the angular position of the r-tick labels
 on a polar plot via :func:`~matplotlib.Axes.axes.set_rlabel_position`.
 
+
 Date handling
 -------------
 

lib/matplotlib/collections.py

@@ -658,6 +658,10 @@ def update_scalarmappable(self):
         elif self._is_stroked:
             self._edgecolors = self.to_rgba(self._A, self._alpha)
 
+    def get_fill(self):
+        'return whether fill is set'
+        return self._is_filled
+
     def update_from(self, other):
         'copy properties from other to self'
 

lib/matplotlib/legend.py

@@ -39,7 +39,7 @@
 from matplotlib.lines import Line2D
 from matplotlib.patches import Patch, Rectangle, Shadow, FancyBboxPatch
 from matplotlib.collections import LineCollection, RegularPolyCollection, \
-     CircleCollection, PathCollection
+     CircleCollection, PathCollection, PolyCollection
 from matplotlib.transforms import Bbox, BboxBase, TransformedBbox
 from matplotlib.transforms import BboxTransformTo, BboxTransformFrom
 
@@ -487,7 +487,8 @@ def _approx_text_height(self, renderer=None):
         BarContainer: legend_handler.HandlerPatch(
                         update_func=legend_handler.update_from_first_child),
         tuple: legend_handler.HandlerTuple(),
-        PathCollection: legend_handler.HandlerPathCollection()
+        PathCollection: legend_handler.HandlerPathCollection(),
+        PolyCollection: legend_handler.HandlerPolyCollection()
         }
 
     # (get|set|update)_default_handler_maps are public interfaces to

lib/matplotlib/legend_handler.py

@@ -35,6 +35,7 @@ def legend_artist(self, legend, orig_handle, fontsize, handlebox):
 from matplotlib.lines import Line2D
 from matplotlib.patches import Rectangle
 import matplotlib.collections as mcoll
+import matplotlib.colors as mcolors
 
 
 def update_from_first_child(tgt, src):
@@ -581,3 +582,38 @@ def create_artists(self, legend, orig_handle,
             a_list.extend(_a_list)
 
         return a_list
+
+
+class HandlerPolyCollection(HandlerBase):
+    """
+    Handler for PolyCollection used in fill_between and stackplot.
+    """
+    def _update_prop(self, legend_handle, orig_handle):
+        def first_color(colors):
+            colors = mcolors.colorConverter.to_rgba_array(colors)
+            if len(colors):
+                return colors[0]
+            else:
+                return "none"
+        def get_first(prop_array):
+            if len(prop_array):
+                return prop_array[0]
+            else:
+                return None
+        legend_handle.set_edgecolor(first_color(orig_handle.get_edgecolor()))
+        legend_handle.set_facecolor(first_color(orig_handle.get_facecolor()))
+        legend_handle.set_fill(orig_handle.get_fill())
+        legend_handle.set_hatch(orig_handle.get_hatch())
+        legend_handle.set_linewidth(get_first(orig_handle.get_linewidths()))
+        legend_handle.set_linestyle(get_first(orig_handle.get_linestyles()))
+        legend_handle.set_transform(get_first(orig_handle.get_transforms()))
+        legend_handle.set_figure(orig_handle.get_figure())
+        legend_handle.set_alpha(orig_handle.get_alpha())
+
+    def create_artists(self, legend, orig_handle,
+                       xdescent, ydescent, width, height, fontsize, trans):
+        p = Rectangle(xy=(-xdescent, -ydescent),
+                      width=width, height=height)
+        self.update_prop(p, orig_handle, legend)
+        p.set_transform(trans)
+        return [p]

lib/matplotlib/stackplot.py

@@ -41,6 +41,9 @@ def stackplot(axes, x, *args, **kwargs):
                 can be found at http://www.leebyron.com/else/streamgraph/.
 
 
+    *labels* : A list or tuple of labels to assign to each data series.
+
+
     *colors* : A list or tuple of colors. These will be cycled through and
                used to colour the stacked areas.
                All other keyword arguments are passed to
@@ -49,18 +52,15 @@ def stackplot(axes, x, *args, **kwargs):
     Returns *r* : A list of
     :class:`~matplotlib.collections.PolyCollection`, one for each
     element in the stacked area plot.
-
-    Note that :class:`~matplotlib.legend.Legend` does not support
-    :class:`~matplotlib.collections.PolyCollection` objects.  To create a
-    legend on a stackplot, use a proxy artist:
-    http://matplotlib.org/users/legend_guide.html#using-proxy-artist
     """
 
     if len(args) == 1:
         y = np.atleast_2d(*args)
     elif len(args) > 1:
         y = np.row_stack(args)
 
+    labels = iter(kwargs.pop('labels', []))
+
     colors = kwargs.pop('colors', None)
     if colors is not None:
         axes.set_color_cycle(colors)
@@ -104,12 +104,14 @@ def stackplot(axes, x, *args, **kwargs):
     # Color between x = 0 and the first array.
     r.append(axes.fill_between(x, first_line, stack[0, :],
                                facecolor=six.next(axes._get_lines.color_cycle),
+                               label= six.next(labels, None),
                                **kwargs))
 
     # Color between array i-1 and array i
     for i in xrange(len(y) - 1):
         color = six.next(axes._get_lines.color_cycle)
         r.append(axes.fill_between(x, stack[i, :], stack[i + 1, :],
                                    facecolor= color,
+                                   label= six.next(labels, None),
                                    **kwargs))
     return r

lib/matplotlib/tests/baseline_images/test_legend/legend_stackplot.png

@@ -218,6 +218,22 @@ def test_legend_handler_map(self):
         handles_labels.assert_called_with({'1': 2})
 
 
+@image_comparison(baseline_images=['legend_stackplot'], extensions=['png'])
+def test_legend_stackplot():
+    '''test legend for PolyCollection using stackplot'''
+    # related to #1341, #1943, and PR #3303
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    x = np.linspace(0, 10, 10)
+    y1 = 1.0 * x
+    y2 = 2.0 * x + 1
+    y3 = 3.0 * x + 2
+    ax.stackplot(x, y1, y2, y3, labels=['y1', 'y2', 'y3'])
+    ax.set_xlim((0, 10))
+    ax.set_ylim((0, 70))
+    ax.legend(loc=0)
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=['-s', '--with-doctest'], exit=False)