Merge pull request #5617 from three-comrades/legend-tuple-handler-improve

WeatherGod · WeatherGod · commit 4230879c5019 · 2016-03-19T17:28:52.000-04:00
Legend tuple handler improve
diff --git a/doc/users/legend_guide.rst b/doc/users/legend_guide.rst
@@ -198,6 +198,21 @@ following example demonstrates combining two legend keys on top of one another:
 
     plt.legend([red_dot, (red_dot, white_cross)], ["Attr A", "Attr A+B"])
 
+The :class:`~matplotlib.legend_handler.HandlerTuple` class can also be used to
+assign several legend keys to the same entry:
+
+.. plot::
+    :include-source:
+
+    import matplotlib.pyplot as plt
+    from matplotlib.legend_handler import HandlerLine2D, HandlerTuple
+
+    p1, = plt.plot([1, 2.5, 3], 'r-d')
+    p2, = plt.plot([3, 2, 1], 'k-o')
+
+    l = plt.legend([(p1, p2)], ['Two keys'], numpoints=1,
+                   handler_map={tuple: HandlerTuple(ndivide=None)})
+
 
 Implementing a custom legend handler
 ------------------------------------
diff --git a/doc/users/whats_new/multiple_legend_keys.rst b/doc/users/whats_new/multiple_legend_keys.rst
@@ -0,0 +1,10 @@
+Multiple legend keys for legend entries
+---------------------------------------
+
+A legend entry can now contain more than one legend key. The extended
+``HandlerTuple`` class now accepts two parameters: ``ndivide`` divides the
+legend area in the specified number of sections; ``pad`` changes the padding
+between the legend keys.
+
+.. plot:: mpl_examples/pylab_examples/legend_demo6.py
+
diff --git a/examples/pylab_examples/legend_demo6.py b/examples/pylab_examples/legend_demo6.py
@@ -0,0 +1,35 @@
+"""
+Showcases legend entries with more than one legend key.
+"""
+import matplotlib.pyplot as plt
+from matplotlib.legend_handler import HandlerTuple
+
+fig, (ax1, ax2) = plt.subplots(2, 1)
+
+# First plot: two legend keys for a single entry
+p1 = ax1.scatter([1], [5], c='r', marker='s', s=100)
+p2 = ax1.scatter([3], [2], c='b', marker='o', s=100)
+# `plot` returns a list, but we want the handle - thus the comma on the left
+p3, = ax1.plot([1, 5], [4, 4], 'm-d')
+
+# Assign two of the handles to the same legend entry by putting them in a tuple
+# and using a generic handler map (which would be used for any additional
+# tuples of handles like (p1, p3)).
+l = ax1.legend([(p1, p3), p2], ['two keys', 'one key'], scatterpoints=1,
+               numpoints=1, handler_map={tuple: HandlerTuple(ndivide=None)})
+
+# Second plot: plot two bar charts on top of each other and change the padding
+# between the legend keys
+x_left = [1, 2, 3]
+y_pos = [1, 3, 2]
+y_neg = [2, 1, 4]
+
+rneg = ax2.bar(x_left, y_neg, width=0.5, color='w', hatch='///', label='-1')
+rpos = ax2.bar(x_left, y_pos, width=0.5, color='k', label='+1')
+
+# Treat each legend entry differently by using specific `HandlerTuple`s
+l = ax2.legend([(rpos, rneg), (rneg, rpos)], ['pad!=0', 'pad=0'],
+               handler_map={(rpos, rneg): HandlerTuple(ndivide=None),
+                            (rneg, rpos): HandlerTuple(ndivide=None, pad=0.)})
+
+plt.show()
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -29,6 +29,7 @@ def legend_artist(self, legend, orig_handle, fontsize, handlebox):
 
 from matplotlib.externals import six
 from matplotlib.externals.six.moves import zip
+from itertools import cycle
 
 import numpy as np
 
@@ -150,13 +151,14 @@ def get_xdata(self, legend, xdescent, ydescent, width, height, fontsize):
         if numpoints > 1:
             # we put some pad here to compensate the size of the
             # marker
-            xdata = np.linspace(-xdescent + self._marker_pad * fontsize,
-                                width - self._marker_pad * fontsize,
+            pad = self._marker_pad * fontsize
+            xdata = np.linspace(-xdescent + pad,
+                                -xdescent + width - pad,
                                 numpoints)
             xdata_marker = xdata
         elif numpoints == 1:
-            xdata = np.linspace(-xdescent, width, 2)
-            xdata_marker = [0.5 * width - 0.5 * xdescent]
+            xdata = np.linspace(-xdescent, -xdescent+width, 2)
+            xdata_marker = [-xdescent + 0.5 * width]
 
         return xdata, xdata_marker
 
@@ -499,6 +501,7 @@ def create_artists(self, legend, orig_handle,
 
         return artists
 
+
 class HandlerStem(HandlerNpointsYoffsets):
     """
     Handler for Errorbars
@@ -565,21 +568,60 @@ def create_artists(self, legend, orig_handle,
 
 class HandlerTuple(HandlerBase):
     """
-    Handler for Tuple
+    Handler for Tuple.
+
+    Additional kwargs are passed through to `HandlerBase`.
+
+    Parameters
+    ----------
+
+    ndivide : int, optional
+        The number of sections to divide the legend area into.  If None,
+        use the length of the input tuple. Default is 1.
+
+
+    pad : float, optional
+        If None, fall back to `legend.borderpad` as the default.
+        In units of fraction of font size. Default is None.
+
+
+
     """
-    def __init__(self, **kwargs):
+    def __init__(self, ndivide=1, pad=None, **kwargs):
+
+        self._ndivide = ndivide
+        self._pad = pad
         HandlerBase.__init__(self, **kwargs)
 
     def create_artists(self, legend, orig_handle,
                        xdescent, ydescent, width, height, fontsize,
                        trans):
 
         handler_map = legend.get_legend_handler_map()
+
+        if self._ndivide is None:
+            ndivide = len(orig_handle)
+        else:
+            ndivide = self._ndivide
+
+        if self._pad is None:
+            pad = legend.borderpad * fontsize
+        else:
+            pad = self._pad * fontsize
+
+        if ndivide > 1:
+            width = (width - pad*(ndivide - 1)) / ndivide
+
+        xds = [xdescent - (width + pad) * i for i in range(ndivide)]
+        xds_cycle = cycle(xds)
+
         a_list = []
         for handle1 in orig_handle:
             handler = legend.get_legend_handler(handler_map, handle1)
             _a_list = handler.create_artists(legend, handle1,
-                                             xdescent, ydescent, width, height,
+                                             six.next(xds_cycle),
+                                             ydescent,
+                                             width, height,
                                              fontsize,
                                              trans)
             a_list.extend(_a_list)
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/legend_multiple_keys.png b/lib/matplotlib/tests/baseline_images/test_legend/legend_multiple_keys.png
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -18,6 +18,7 @@
 import matplotlib.patches as mpatches
 import matplotlib.transforms as mtrans
 
+from matplotlib.legend_handler import HandlerTuple
 
 @image_comparison(baseline_images=['legend_auto1'], remove_text=True)
 def test_legend_auto1():
@@ -77,6 +78,21 @@ def test_labels_first():
     ax.legend(loc=0, markerfirst=False)
 
 
+@image_comparison(baseline_images=['legend_multiple_keys'], extensions=['png'],
+                  remove_text=True)
+def test_multiple_keys():
+    # test legend entries with multiple keys
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    p1, = ax.plot([1, 2, 3], '-o')
+    p2, = ax.plot([2, 3, 4], '-x')
+    p3, = ax.plot([3, 4, 5], '-d')
+    ax.legend([(p1, p2), (p2, p1), p3], ['two keys', 'pad=0', 'one key'],
+              numpoints=1,
+              handler_map={(p1, p2): HandlerTuple(ndivide=None),
+                           (p2, p1): HandlerTuple(ndivide=None, pad=0)})
+
+
 @image_comparison(baseline_images=['rgba_alpha'],
                   extensions=['png'], remove_text=True)
 def test_alpha_rgba():
