adding legend handler for PolyCollection

Francis Colas · Francis Colas · commit e6958c3c6e1b · 2014-10-21T23:06:57.000+02:00
diff --git a/lib/matplotlib/collections.py b/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'
 
diff --git a/lib/matplotlib/legend.py b/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
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -35,7 +35,8 @@ 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
+import matplotlib.artist as artist
 
 def update_from_first_child(tgt, src):
     tgt.update_from(src.get_children()[0])
@@ -581,3 +582,39 @@ 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
+        #artist.Artist.update_from(legend_handle, orig_handle)
+        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]
