add cache_frame_data kwarg into FuncAnimation. fixes #8528.

pclove1 · pclove1 · commit 017d88b493ee · 2018-09-03T15:04:28.000-07:00
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -1621,9 +1621,13 @@ def init_func() -> iterable_of_artists:
        However, they will be drawn on top of any previous artists, regardless
        of their zorder.  Defaults to ``False``.
 
+    cache_frame_data : bool, optional
+       Controls whether frame data is cached. Defaults to ``True``.
+       Disabling cache might be helpful when frames contain large objects.
+
     '''
     def __init__(self, fig, func, frames=None, init_func=None, fargs=None,
-                 save_count=None, **kwargs):
+                 save_count=None, cache_frame_data=True, **kwargs):
         if fargs:
             self._args = fargs
         else:
@@ -1671,6 +1675,8 @@ def __init__(self, fig, func, frames=None, init_func=None, fargs=None,
         # for a single frame from init, which is not what we want.
         self._save_seq = []
 
+        self._cache_frame_data = cache_frame_data
+
     def new_frame_seq(self):
         # Use the generating function to generate a new frame sequence
         return self._iter_gen()
@@ -1725,8 +1731,9 @@ def _init_draw(self):
         self._save_seq = []
 
     def _draw_frame(self, framedata):
-        # Save the data for potential saving of movies.
-        self._save_seq.append(framedata)
+        if self._cache_frame_data:
+            # Save the data for potential saving of movies.
+            self._save_seq.append(framedata)
 
         # Make sure to respect save_count (keep only the last save_count
         # around)
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -1,6 +1,7 @@
 import os
 from pathlib import Path
 import sys
+import weakref
 
 import numpy as np
 import pytest
@@ -260,3 +261,86 @@ def test_failing_ffmpeg(tmpdir, monkeypatch):
                 make_animation().save("test.mpeg")
     finally:
         animation.writers.reset_available_writers()
+
+
+def test_funcanimation_holds_frames_by_default():
+    fig, ax = plt.subplots()
+    line, = ax.plot([], [])
+
+    class Frame(dict):
+        # this subclassing enables to use weakref.ref()
+        pass
+
+    def init():
+        line.set_data([], [])
+        return line,
+
+    def animate(frame):
+        line.set_data(frame['x'], frame['y'])
+        return line,
+
+    frames_generated = []
+
+    def frames_generator():
+        for _ in range(5):
+            x = np.linspace(0, 10, 100)
+            y = np.random.rand(100)
+
+            frame = Frame(x=x, y=y)
+
+            # collect weak references to frames
+            # to validate their references later
+            frames_generated.append(weakref.ref(frame))
+
+            yield frame
+
+    anim = animation.FuncAnimation(fig, animate, init_func=init,
+                                   frames=frames_generator, save_count=5)
+
+    writer = NullMovieWriter()
+    anim.save('unused.null', writer=writer)
+
+    for f in frames_generated:
+        assert f() is not None
+
+
+def test_funcanimation_does_not_hold_frames_when_cache_frame_data_is_False():
+    fig, ax = plt.subplots()
+    line, = ax.plot([], [])
+
+    class Frame(dict):
+        # this subclassing enables to use weakref.ref()
+        pass
+
+    def init():
+        line.set_data([], [])
+        return line,
+
+    def animate(frame):
+        line.set_data(frame['x'], frame['y'])
+        return line,
+
+    frames_generated = []
+
+    def frames_generator():
+        for _ in range(5):
+            x = np.linspace(0, 10, 100)
+            y = np.random.rand(100)
+
+            frame = Frame(x=x, y=y)
+
+            # collect weak references to frames
+            # to validate their references later
+            frames_generated.append(weakref.ref(frame))
+
+            yield frame
+
+    anim = animation.FuncAnimation(fig, animate, init_func=init,
+                                   frames=frames_generator, save_count=5,
+                                   cache_frame_data=False)
+
+    writer = NullMovieWriter()
+    anim.save('unused.null', writer=writer)
+
+    for f in frames_generated:
+        assert f() is None
