raise RuntimeError appropriately for animation update func

johnnyEmpires · johnnyEmpires · commit 2c15ab4668dc · 2020-06-20T22:15:29.000+09:00
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -1716,9 +1716,19 @@ def _draw_frame(self, framedata):
         # 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
@@ -272,3 +272,40 @@ 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 init():
+        pass
+
+    def animate_case_1(i):
+        # user forgot to return (returns None)
+        line.set_data([0, 1], [0, i])
+        # return line
+
+    def animate_case_2(i):
+        # user forgot to put comma or return a sequence
+        # TypeError will be raised (same with returning a number or bool)
+        line.set_data([0, 1], [0, i])
+        return line
+
+    def animate_case_3(i):
+        # user (for some reason) returned a string...AttributeError is raised
+        line.set_data([0, 1], [0, i])
+        return 'a string'
+
+    def animate_case_4(i):
+        # user returns a sequence other objects instead of Artist.
+        line.set_data([0, 1], [0, i])
+        return 'a string',
+
+    with pytest.raises(RuntimeError) as context:
+        animation.FuncAnimation(fig, animate_case_1, blit=True)
+        animation.FuncAnimation(fig, animate_case_2, blit=True)
+        animation.FuncAnimation(fig, animate_case_3, blit=True)
+        animation.FuncAnimation(fig, animate_case_4, blit=True)
