Merge pull request #17699 from johnnyEmpires/draw-frame-error-check

jklymak · web-flow · commit 4723dabd7885 · 2020-07-18T13:56:56.000-07:00
raise RuntimeError appropriately for animation update func
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -1716,10 +1716,22 @@ def _draw_frame(self, framedata):
         # Call the func with framedata and args. If blitting is desired,
         # func needs to return a sequence of any artists that were modified.
         self._drawn_artists = self._func(framedata, *self._args)
+
         if self._blit:
-            if self._drawn_artists is None:
-                raise RuntimeError('The animation function must return a '
+
+            err = RuntimeError('The animation function must return a '
                                    'sequence of Artist objects.')
+            try:
+                # check if a sequence
+                iter(self._drawn_artists)
+            except TypeError:
+                raise err
+
+            # check each item if is artist
+            for i in self._drawn_artists:
+                if not isinstance(i, mpl.artist.Artist):
+                    raise err
+
             self._drawn_artists = sorted(self._drawn_artists,
                                          key=lambda x: x.get_zorder())
 
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -10,6 +10,7 @@
 import matplotlib as mpl
 from matplotlib import pyplot as plt
 from matplotlib import animation
+from matplotlib import artist
 
 
 class NullMovieWriter(animation.AbstractMovieWriter):
@@ -272,3 +273,36 @@ def frames_generator():
         # If cache_frame_data is True, then the weakref should be alive;
         # if cache_frame_data is False, then the weakref should be dead (None).
         assert (f() is None) != cache_frame_data
+
+
+def test_draw_frame():
+    # test _draw_frame method
+
+    fig, ax = plt.subplots()
+    line, = ax.plot([])
+
+    def animate(i, arg):
+        # general update func
+        line.set_data([0, 1], [0, i])
+        if arg:
+            return arg
+
+    with pytest.raises(RuntimeError):
+
+        # user forgot to return (returns None)
+        animation.FuncAnimation(fig, animate, blit=True, fargs=(None,))
+
+        # user (for some reason) returned a string...AttributeError is raised
+        animation.FuncAnimation(fig, animate, blit=True, fargs=('string', ))
+
+        # user (for some reason) returned a string...AttributeError is raised
+        animation.FuncAnimation(fig, animate, blit=True, fargs=(1, ))
+
+        # user returns a sequence of other objects
+        # e.g. a string instead of Artist
+        animation.FuncAnimation(fig, animate, blit=True, fargs=(('string',), ))
+
+        # user forgot to put comma or return a sequence
+        # TypeError will be raised (same with returning a number or bool)
+        artist_obj = artist.Artist()
+        animation.FuncAnimation(fig, animate, blit=True, fargs=(artist_obj, ))
