test_animation.py 8.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274
  1. import os
  2. from pathlib import Path
  3. import subprocess
  4. import sys
  5. import weakref
  6. import numpy as np
  7. import pytest
  8. import matplotlib as mpl
  9. from matplotlib import pyplot as plt
  10. from matplotlib import animation
  11. class NullMovieWriter(animation.AbstractMovieWriter):
  12. """
  13. A minimal MovieWriter. It doesn't actually write anything.
  14. It just saves the arguments that were given to the setup() and
  15. grab_frame() methods as attributes, and counts how many times
  16. grab_frame() is called.
  17. This class doesn't have an __init__ method with the appropriate
  18. signature, and it doesn't define an isAvailable() method, so
  19. it cannot be added to the 'writers' registry.
  20. """
  21. def setup(self, fig, outfile, dpi, *args):
  22. self.fig = fig
  23. self.outfile = outfile
  24. self.dpi = dpi
  25. self.args = args
  26. self._count = 0
  27. def grab_frame(self, **savefig_kwargs):
  28. self.savefig_kwargs = savefig_kwargs
  29. self._count += 1
  30. def finish(self):
  31. pass
  32. def make_animation(**kwargs):
  33. fig, ax = plt.subplots()
  34. line, = ax.plot([])
  35. def init():
  36. pass
  37. def animate(i):
  38. line.set_data([0, 1], [0, i])
  39. return line,
  40. return animation.FuncAnimation(fig, animate, **kwargs)
  41. def test_null_movie_writer():
  42. # Test running an animation with NullMovieWriter.
  43. num_frames = 5
  44. anim = make_animation(frames=num_frames)
  45. filename = "unused.null"
  46. dpi = 50
  47. savefig_kwargs = dict(foo=0)
  48. writer = NullMovieWriter()
  49. anim.save(filename, dpi=dpi, writer=writer,
  50. savefig_kwargs=savefig_kwargs)
  51. assert writer.fig == plt.figure(1) # The figure used by make_animation.
  52. assert writer.outfile == filename
  53. assert writer.dpi == dpi
  54. assert writer.args == ()
  55. assert writer.savefig_kwargs == savefig_kwargs
  56. assert writer._count == num_frames
  57. def test_movie_writer_dpi_default():
  58. class DummyMovieWriter(animation.MovieWriter):
  59. def _run(self):
  60. pass
  61. # Test setting up movie writer with figure.dpi default.
  62. fig = plt.figure()
  63. filename = "unused.null"
  64. fps = 5
  65. codec = "unused"
  66. bitrate = 1
  67. extra_args = ["unused"]
  68. writer = DummyMovieWriter(fps, codec, bitrate, extra_args)
  69. writer.setup(fig, filename)
  70. assert writer.dpi == fig.dpi
  71. @animation.writers.register('null')
  72. class RegisteredNullMovieWriter(NullMovieWriter):
  73. # To be able to add NullMovieWriter to the 'writers' registry,
  74. # we must define an __init__ method with a specific signature,
  75. # and we must define the class method isAvailable().
  76. # (These methods are not actually required to use an instance
  77. # of this class as the 'writer' argument of Animation.save().)
  78. def __init__(self, fps=None, codec=None, bitrate=None,
  79. extra_args=None, metadata=None):
  80. pass
  81. @classmethod
  82. def isAvailable(cls):
  83. return True
  84. WRITER_OUTPUT = [
  85. ('ffmpeg', 'movie.mp4'),
  86. ('ffmpeg_file', 'movie.mp4'),
  87. ('avconv', 'movie.mp4'),
  88. ('avconv_file', 'movie.mp4'),
  89. ('imagemagick', 'movie.gif'),
  90. ('imagemagick_file', 'movie.gif'),
  91. ('pillow', 'movie.gif'),
  92. ('html', 'movie.html'),
  93. ('null', 'movie.null')
  94. ]
  95. WRITER_OUTPUT += [
  96. (writer, Path(output)) for writer, output in WRITER_OUTPUT]
  97. # Smoke test for saving animations. In the future, we should probably
  98. # design more sophisticated tests which compare resulting frames a-la
  99. # matplotlib.testing.image_comparison
  100. @pytest.mark.parametrize('writer, output', WRITER_OUTPUT)
  101. def test_save_animation_smoketest(tmpdir, writer, output):
  102. if not animation.writers.is_available(writer):
  103. pytest.skip("writer '%s' not available on this system" % writer)
  104. fig, ax = plt.subplots()
  105. line, = ax.plot([], [])
  106. ax.set_xlim(0, 10)
  107. ax.set_ylim(-1, 1)
  108. dpi = None
  109. codec = None
  110. if writer == 'ffmpeg':
  111. # Issue #8253
  112. fig.set_size_inches((10.85, 9.21))
  113. dpi = 100.
  114. codec = 'h264'
  115. def init():
  116. line.set_data([], [])
  117. return line,
  118. def animate(i):
  119. x = np.linspace(0, 10, 100)
  120. y = np.sin(x + i)
  121. line.set_data(x, y)
  122. return line,
  123. # Use temporary directory for the file-based writers, which produce a file
  124. # per frame with known names.
  125. with tmpdir.as_cwd():
  126. anim = animation.FuncAnimation(fig, animate, init_func=init, frames=5)
  127. anim.save(output, fps=30, writer=writer, bitrate=500, dpi=dpi,
  128. codec=codec)
  129. def test_no_length_frames():
  130. (make_animation(frames=iter(range(5)))
  131. .save('unused.null', writer=NullMovieWriter()))
  132. def test_movie_writer_registry():
  133. assert len(animation.writers._registered) > 0
  134. mpl.rcParams['animation.ffmpeg_path'] = "not_available_ever_xxxx"
  135. assert not animation.writers.is_available("ffmpeg")
  136. # something guaranteed to be available in path and exits immediately
  137. bin = "true" if sys.platform != 'win32' else "where"
  138. mpl.rcParams['animation.ffmpeg_path'] = bin
  139. assert animation.writers.is_available("ffmpeg")
  140. @pytest.mark.parametrize(
  141. "method_name",
  142. [pytest.param("to_html5_video", marks=pytest.mark.skipif(
  143. not animation.writers.is_available(mpl.rcParams["animation.writer"]),
  144. reason="animation writer not installed")),
  145. "to_jshtml"])
  146. def test_embed_limit(method_name, caplog, tmpdir):
  147. caplog.set_level("WARNING")
  148. with tmpdir.as_cwd():
  149. with mpl.rc_context({"animation.embed_limit": 1e-6}): # ~1 byte.
  150. getattr(make_animation(frames=1), method_name)()
  151. assert len(caplog.records) == 1
  152. record, = caplog.records
  153. assert (record.name == "matplotlib.animation"
  154. and record.levelname == "WARNING")
  155. @pytest.mark.parametrize(
  156. "method_name",
  157. [pytest.param("to_html5_video", marks=pytest.mark.skipif(
  158. not animation.writers.is_available(mpl.rcParams["animation.writer"]),
  159. reason="animation writer not installed")),
  160. "to_jshtml"])
  161. def test_cleanup_temporaries(method_name, tmpdir):
  162. with tmpdir.as_cwd():
  163. getattr(make_animation(frames=1), method_name)()
  164. assert list(Path(str(tmpdir)).iterdir()) == []
  165. @pytest.mark.skipif(os.name != "posix", reason="requires a POSIX OS")
  166. def test_failing_ffmpeg(tmpdir, monkeypatch):
  167. """
  168. Test that we correctly raise a CalledProcessError when ffmpeg fails.
  169. To do so, mock ffmpeg using a simple executable shell script that
  170. succeeds when called with no arguments (so that it gets registered by
  171. `isAvailable`), but fails otherwise, and add it to the $PATH.
  172. """
  173. with tmpdir.as_cwd():
  174. monkeypatch.setenv("PATH", ".:" + os.environ["PATH"])
  175. exe_path = Path(str(tmpdir), "ffmpeg")
  176. exe_path.write_text("#!/bin/sh\n"
  177. "[[ $@ -eq 0 ]]\n")
  178. os.chmod(str(exe_path), 0o755)
  179. with pytest.raises(subprocess.CalledProcessError):
  180. make_animation().save("test.mpeg")
  181. @pytest.mark.parametrize("cache_frame_data", [False, True])
  182. def test_funcanimation_cache_frame_data(cache_frame_data):
  183. fig, ax = plt.subplots()
  184. line, = ax.plot([], [])
  185. class Frame(dict):
  186. # this subclassing enables to use weakref.ref()
  187. pass
  188. def init():
  189. line.set_data([], [])
  190. return line,
  191. def animate(frame):
  192. line.set_data(frame['x'], frame['y'])
  193. return line,
  194. frames_generated = []
  195. def frames_generator():
  196. for _ in range(5):
  197. x = np.linspace(0, 10, 100)
  198. y = np.random.rand(100)
  199. frame = Frame(x=x, y=y)
  200. # collect weak references to frames
  201. # to validate their references later
  202. frames_generated.append(weakref.ref(frame))
  203. yield frame
  204. anim = animation.FuncAnimation(fig, animate, init_func=init,
  205. frames=frames_generator,
  206. cache_frame_data=cache_frame_data)
  207. writer = NullMovieWriter()
  208. anim.save('unused.null', writer=writer)
  209. assert len(frames_generated) == 5
  210. for f in frames_generated:
  211. # If cache_frame_data is True, then the weakref should be alive;
  212. # if cache_frame_data is False, then the weakref should be dead (None).
  213. assert (f() is None) != cache_frame_data