FIX: support np.flatiter as input again

tacaswell · tacaswell · commit 8d032996035b · 2016-07-11T22:31:09.000-05:00
This was unintentionally broken by bdaaf59 which went in as part of PR #5556 which was fixing support for pandas series.
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -2544,6 +2544,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)
