Merge pull request #1536 from jakevdp/anim_avconv

dopplershift · dopplershift · commit 7924451674e2 · 2012-12-06T20:46:20.000-08:00
ENH: add AVConv movie writer for animations
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -53,6 +53,9 @@ def list(self):
         ''' Get a list of available MovieWriters.'''
         return self.avail.keys()
 
+    def is_available(self, name):
+        return name in self.avail
+
     def __getitem__(self, name):
         if not self.avail:
             raise RuntimeError("No MovieWriters available!")
@@ -390,6 +393,24 @@ def _args(self):
                 self._base_temp_name()] + self.output_args
 
 
+# Base class of avconv information.  AVConv has identical arguments to
+# FFMpeg
+class AVConvBase(FFMpegBase):
+    exec_key = 'animation.avconv_path'
+    args_key = 'animation.avconv_args' 
+
+
+# Combine AVConv options with pipe-based writing
+@writers.register('avconv')
+class AVConvWriter(AVConvBase, FFMpegWriter):
+    pass
+
+# Combine AVConv options with file-based writing
+@writers.register('avconv_file')
+class AVConvFileWriter(AVConvBase, FFMpegFileWriter):
+    pass
+
+
 # Base class of mencoder information. Contains configuration key information
 # as well as arguments for controlling *output*
 class MencoderBase:
diff --git a/lib/matplotlib/rcsetup.py b/lib/matplotlib/rcsetup.py
@@ -336,7 +336,9 @@ def validate_hinting(s):
                                            ['xelatex', 'lualatex', 'pdflatex'])
 
 validate_movie_writer = ValidateInStrings('animation.writer',
-    ['ffmpeg', 'ffmpeg_file', 'mencoder', 'mencoder_file',
+    ['ffmpeg', 'ffmpeg_file',
+     'avconv', 'avconv_file',
+     'mencoder', 'mencoder_file',
      'imagemagick', 'imagemagick_file'])
 
 validate_movie_frame_fmt = ValidateInStrings('animation.frame_format',
@@ -636,7 +638,9 @@ def __call__(self, s):
     'animation.frame_format' : ['png', validate_movie_frame_fmt], # Controls image format when frames are written to disk
     'animation.ffmpeg_path' : ['ffmpeg', str], # Path to FFMPEG binary. If just binary name, subprocess uses $PATH.
     'animation.ffmpeg_args' : ['', validate_stringlist], # Additional arguments for ffmpeg movie writer (using pipes)
-    'animation.mencoder_path' : ['mencoder', str], # Path to FFMPEG binary. If just binary name, subprocess uses $PATH.
+    'animation.avconv_path' : ['avconv', str], # Path to AVConv binary. If just binary name, subprocess uses $PATH.
+    'animation.avconv_args' : ['', validate_stringlist], # Additional arguments for avconv movie writer (using pipes)
+    'animation.mencoder_path' : ['mencoder', str], # Path to MENCODER binary. If just binary name, subprocess uses $PATH.
     'animation.mencoder_args' : ['', validate_stringlist], # Additional arguments for mencoder movie writer (using pipes)
     'animation.convert_path' : ['convert', str], # Path to convert binary. If just binary name, subprocess uses $PATH
     'animation.convert_args' : ['', validate_stringlist], # Additional arguments for mencoder movie writer (using pipes)
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -0,0 +1,49 @@
+import os
+import tempfile
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib import animation
+from matplotlib.testing.noseclasses import KnownFailureTest
+
+
+WRITER_OUTPUT = dict(ffmpeg='mp4', ffmpeg_file='mp4',
+                     mencoder='mp4', mencoder_file='mp4',
+                     avconv='mp4', avconv_file='mp4',
+                     imagemagick='gif', imagemagick_file='gif')
+                     
+
+
+# Smoke test for saving animations.  In the future, we should probably
+# design more sophisticated tests which compare resulting frames a-la
+# matplotlib.testing.image_comparison
+def test_save_animation_smoketest():
+    for writer, extension in WRITER_OUTPUT.iteritems():
+        yield check_save_animation, writer, extension
+
+
+def check_save_animation(writer, extension='mp4'):
+    if not animation.writers.is_available(writer):
+        raise KnownFailureTest("writer '%s' not available on this system"
+                               % writer)
+    fig, ax = plt.subplots()
+    line, = ax.plot([], [])
+
+    def init():
+        line.set_data([], [])
+        return line,
+
+    def animate(i):
+        x = np.linspace(0, 10, 100)
+        y = np.sin(x + i)
+        line.set_data(x, y)
+        return line,
+
+    # Use NamedTemporaryFile: will be automatically deleted
+    F = tempfile.NamedTemporaryFile(suffix='.' + extension)
+    anim = animation.FuncAnimation(fig, animate, init_func=init, frames=5)
+    anim.save(F.name, fps=30, writer=writer)
+
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule()
diff --git a/matplotlibrc.template b/matplotlibrc.template
@@ -434,7 +434,11 @@ text.hinting_factor : 8 # Specifies the amount of softness for hinting in the
 #animation.frame_format: 'png'     # Controls frame format used by temp files
 #animation.ffmpeg_path: 'ffmpeg'   # Path to ffmpeg binary. Without full path
                                    # $PATH is searched
-#animation.ffmpeg_args: ''         # Additional arugments to pass to mencoder
-#animation.mencoder_path: 'ffmpeg' # Path to mencoder binary. Without full path
+#animation.ffmpeg_args: ''         # Additional arguments to pass to ffmpeg
+#animation.avconv_path: 'avconv'   # Path to avconv binary. Without full path
                                    # $PATH is searched
-#animation.mencoder_args: ''       # Additional arugments to pass to mencoder
+#animation.avconv_args: ''         # Additional arguments to pass to avconv
+#animation.mencoder_path: 'mencoder' 
+			           # Path to mencoder binary. Without full path
+                                   # $PATH is searched
+#animation.mencoder_args: ''       # Additional arguments to pass to mencoder
