Merge pull request #19056 from QuLogic/ffmpeg-file-raw

anntzer · web-flow · commit ecba1d24c90c · 2020-12-18T14:07:23.000+01:00
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -628,15 +628,28 @@ class FFMpegFileWriter(FFMpegBase, FileMovieWriter):
     Frames are written to temporary files on disk and then stitched
     together at the end.
     """
-    supported_formats = ['png', 'jpeg', 'ppm', 'tiff', 'sgi', 'bmp',
-                         'pbm', 'raw', 'rgba']
+    supported_formats = ['png', 'jpeg', 'tiff', 'raw', 'rgba']
 
     def _args(self):
         # Returns the command line parameters for subprocess to use
         # ffmpeg to create a movie using a collection of temp images
-        return [self.bin_path(), '-r', str(self.fps),
-                '-i', self._base_temp_name(),
-                '-vframes', str(self._frame_counter)] + self.output_args
+        args = []
+        # For raw frames, we need to explicitly tell ffmpeg the metadata.
+        if self.frame_format in {'raw', 'rgba'}:
+            args += [
+                '-f', 'image2', '-vcodec', 'rawvideo',
+                '-video_size', '%dx%d' % self.frame_size,
+                '-pixel_format', 'rgba',
+                '-framerate', str(self.fps),
+            ]
+        args += ['-r', str(self.fps), '-i', self._base_temp_name(),
+                 '-vframes', str(self._frame_counter)]
+        # Logging is quieted because subprocess.PIPE has limited buffer size.
+        # If you have a lot of frames in your animation and set logging to
+        # DEBUG, you will have a buffer overrun.
+        if _log.getEffectiveLevel() > logging.DEBUG:
+            args += ['-loglevel', 'error']
+        return [self.bin_path(), *args, *self.output_args]
 
 
 # Base class of avconv information.  AVConv has identical arguments to FFMpeg.
@@ -745,8 +758,7 @@ class ImageMagickFileWriter(ImageMagickBase, FileMovieWriter):
     together at the end.
     """
 
-    supported_formats = ['png', 'jpeg', 'ppm', 'tiff', 'sgi', 'bmp',
-                         'pbm', 'raw', 'rgba']
+    supported_formats = ['png', 'jpeg', 'tiff', 'raw', 'rgba']
 
     def _args(self):
         # Force format: ImageMagick does not recognize 'raw'.
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -141,19 +141,31 @@ def isAvailable(cls):
     ('html', 'movie.html'),
     ('null', 'movie.null')
 ]
-WRITER_OUTPUT += [
-    (writer, Path(output)) for writer, output in WRITER_OUTPUT]
+
+
+def gen_writers():
+    for writer, output in WRITER_OUTPUT:
+        if not animation.writers.is_available(writer):
+            mark = pytest.mark.skip(
+                f"writer '{writer}' not available on this system")
+            yield pytest.param(writer, None, output, marks=[mark])
+            yield pytest.param(writer, None, Path(output), marks=[mark])
+            continue
+
+        writer_class = animation.writers[writer]
+        for frame_format in getattr(writer_class, 'supported_formats', [None]):
+            yield writer, frame_format, output
+            yield writer, frame_format, Path(output)
 
 
 # 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
-@pytest.mark.parametrize('writer, output', WRITER_OUTPUT)
+@pytest.mark.parametrize('writer, frame_format, output', gen_writers())
 @pytest.mark.parametrize('anim', [dict(klass=dict)], indirect=['anim'])
-def test_save_animation_smoketest(tmpdir, writer, output, anim):
-    if not animation.writers.is_available(writer):
-        pytest.skip("writer '%s' not available on this system" % writer)
-
+def test_save_animation_smoketest(tmpdir, writer, frame_format, output, anim):
+    if frame_format is not None:
+        plt.rcParams["animation.frame_format"] = frame_format
     anim = animation.FuncAnimation(**anim)
     dpi = None
     codec = None
