Merge pull request #6712 from tacaswell/mnt_fix_npflatiter_as_input

mdboom · web-flow · commit aeed94b8f0b9 · 2016-08-29T15:06:59.000-04:00
FIX: support np.flatiter is input again
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -2550,6 +2550,15 @@ def index_of(y):
 
 def safe_first_element(obj):
     if isinstance(obj, collections.Iterator):
+        # needed to accept `array.flat` as input.
+        # np.flatiter reports as an instance of collections.Iterator
+        # but can still be indexed via [].
+        # This has the side effect of re-setting the iterator, but
+        # that is acceptable.
+        try:
+            return obj[0]
+        except TypeError:
+            pass
         raise RuntimeError("matplotlib does not support generators "
                            "as input")
     return next(iter(obj))
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -499,3 +499,15 @@ class dummy():
     base_set = mapping[ref(objs[0])]
     for o in objs[1:]:
         assert mapping[ref(o)] is base_set
+
+
+def test_flatiter():
+    x = np.arange(5)
+    it = x.flat
+    assert 0 == next(it)
+    assert 1 == next(it)
+    ret = cbook.safe_first_element(it)
+    assert ret == 0
+
+    assert 0 == next(it)
+    assert 1 == next(it)
