Merge pull request #2634 from tacaswell/animation_iter_length

pelson · pelson · commit 397c27ae6c01 · 2014-01-06T02:10:11.000-08:00
address FuncAnimantion trying to take lengths of generators
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -1000,7 +1000,8 @@ def __init__(self, fig, func, frames=None, init_func=None, fargs=None,
             self._iter_gen = frames
         elif iterable(frames):
             self._iter_gen = lambda: iter(frames)
-            self.save_count = len(frames)
+            if hasattr(frames, '__len__'):
+                self.save_count = len(frames)
         else:
             self._iter_gen = lambda: xrange(frames).__iter__()
             self.save_count = frames
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -58,6 +58,25 @@ def animate(i):
                                "see issues #1891 and #2679")
 
 
+@cleanup
+def test_no_length_frames():
+    fig, ax = plt.subplots()
+    line, = ax.plot([], [])
+
+    def init():
+        line.set_data([], [])
+        return line,
+
+    def animate(i):
+        x = np.linspace(0, 10, 100)
+        y = np.sin(x + i)
+        line.set_data(x, y)
+        return line,
+
+    anim = animation.FuncAnimation(fig, animate, init_func=init,
+                                   frames=iter(range(5)))
+
+
 if __name__ == "__main__":
     import nose
     nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
