matplotlib
diff --git a/‎lib/matplotlib/cbook.py
Copy file name to clipboardExpand all lines: lib/matplotlib/cbook.py
+9Lines changed: 9 additions & 0 deletions b/‎lib/matplotlib/cbook.py
Copy file name to clipboardExpand all lines: lib/matplotlib/cbook.py
+9Lines changed: 9 additions & 0 deletions
diff --git a/‎lib/matplotlib/tests/test_cbook.py
Copy file name to clipboardExpand all lines: lib/matplotlib/tests/test_cbook.py
+12Lines changed: 12 additions & 0 deletions b/‎lib/matplotlib/tests/test_cbook.py
Copy file name to clipboardExpand all lines: lib/matplotlib/tests/test_cbook.py
+12Lines changed: 12 additions & 0 deletions
@@ -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))
 
@@ -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)