Merge pull request #18124 from QuLogic/fix-test_draw_frame

tacaswell · web-flow · commit 9fa9efba2c3b · 2020-07-30T10:12:11.000-04:00
FIX: FuncAnimation._draw_frame exception and testing
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -1723,13 +1723,13 @@ def _draw_frame(self, framedata):
 
         if self._blit:
 
-            err = RuntimeError('The animation function must return a '
-                                   'sequence of Artist objects.')
+            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
+                raise err from None
 
             # check each item if is artist
             for i in self._drawn_artists:
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -10,7 +10,6 @@
 import matplotlib as mpl
 from matplotlib import pyplot as plt
 from matplotlib import animation
-from matplotlib import artist
 
 
 class NullMovieWriter(animation.AbstractMovieWriter):
@@ -275,34 +274,32 @@ def frames_generator():
         assert (f() is None) != cache_frame_data
 
 
-def test_draw_frame():
+@pytest.mark.parametrize('return_value', [
+    # User forgot to return (returns None).
+    None,
+    # User returned a string.
+    'string',
+    # User returned an int.
+    1,
+    # User returns a sequence of other objects, e.g., string instead of Artist.
+    ('string', ),
+    # User forgot to return a sequence (handled in `animate` below.)
+    'artist',
+])
+def test_draw_frame(return_value):
     # test _draw_frame method
 
     fig, ax = plt.subplots()
     line, = ax.plot([])
 
-    def animate(i, arg):
+    def animate(i):
         # general update func
         line.set_data([0, 1], [0, i])
-        if arg:
-            return arg
+        if return_value == 'artist':
+            # *not* a sequence
+            return line
+        else:
+            return return_value
 
     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, ))
+        animation.FuncAnimation(fig, animate, blit=True)
