test code clean up

johnnyEmpires · johnnyEmpires · commit 1808472431b5 · 2020-07-16T20:24:19.000+09:00
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -1716,7 +1716,9 @@ 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:
+
             err = RuntimeError('The animation function must return a '
                                    'sequence of Artist objects.')
             try:
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):
@@ -280,32 +281,29 @@ def test_draw_frame():
     fig, ax = plt.subplots()
     line, = ax.plot([])
 
-    def init():
-        pass
+    def animate(i, arg):
+        # general update func
+        line.set_data([0, 1], [0, i])
+        if arg:
+            return arg
+
+    with pytest.raises(RuntimeError):
 
-    def animate_case_1(i):
         # user forgot to return (returns None)
-        line.set_data([0, 1], [0, i])
-        # return line
+        animation.FuncAnimation(fig, animate, blit=True, fargs=(None,))
 
-    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
+        # user (for some reason) returned a string...AttributeError is raised
+        animation.FuncAnimation(fig, animate, blit=True, fargs=('string', ))
 
-    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'
+        animation.FuncAnimation(fig, animate, blit=True, fargs=(1, ))
 
-    def animate_case_4(i):
-        # user returns a sequence other objects instead of Artist.
-        line.set_data([0, 1], [0, i])
-        return 'a string',
+        # user returns a sequence of other objects
+        # e.g. a string instead of Artist
+        animation.FuncAnimation(fig, animate, blit=True, fargs=(('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)
+        # 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, ))
+    
