Merge pull request #10263 from jklymak/fix-legend-ordereddict

tacaswell · web-flow · commit 4f9458d56074 · 2018-01-17T14:12:05.000-05:00
FIX: (re-allow) legend OrderedDict handles and labels...
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -589,14 +589,17 @@ def __init__(self, parent, handles, labels,
             setattr(self, name, value)
         del locals_view
         # trim handles and labels if illegal label...
-        for label, handle in zip(labels[:], handles[:]):
-                if (isinstance(label, six.string_types) and
-                        label.startswith('_')):
-                    warnings.warn('The handle {!r} has a label of {!r} which '
-                                  'cannot be automatically added to the '
-                                  'legend.'.format(handle, label))
-                    labels.remove(label)
-                    handles.remove(handle)
+        _lab, _hand = [], []
+        for label, handle in zip(labels, handles):
+            if (isinstance(label, six.string_types) and
+                    label.startswith('_')):
+                warnings.warn('The handle {!r} has a label of {!r} which '
+                              'cannot be automatically added to the '
+                              'legend.'.format(handle, label))
+            else:
+                _lab.append(label)
+                _hand.append(handle)
+        labels, handles = _lab, _hand
 
         handles = list(handles)
         if len(handles) < 2:
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -5,6 +5,7 @@
     from unittest import mock
 except ImportError:
     import mock
+import collections
 import numpy as np
 import pytest
 
@@ -53,6 +54,23 @@ def test_legend_kwdocstrings():
     assert stleg == stfig
 
 
+def test_legend_ordereddict():
+    # smoketest that ordereddict inputs work...
+
+    X = np.random.randn(10)
+    Y = np.random.randn(10)
+    labels = ['a'] * 5 + ['b'] * 5
+    colors = ['r'] * 5 + ['g'] * 5
+
+    fig, ax = plt.subplots()
+    for x, y, label, color in zip(X, Y, labels, colors):
+        ax.scatter(x, y, label=label, c=color)
+
+    handles, labels = ax.get_legend_handles_labels()
+    legend = collections.OrderedDict(zip(labels, handles))
+    ax.legend(legend.values(), legend.keys(), loc=6, bbox_to_anchor=(1, .5))
+
+
 @image_comparison(baseline_images=['legend_auto1'], remove_text=True)
 def test_legend_auto1():
     'Test automatic legend placement'
