Pass AbstractMovieWriter.setup kwargs through Animation.save

Author: fredrik-1

lib/matplotlib/animation.py

@@ -210,7 +210,7 @@ def saving(self, fig, outfile, dpi, *args, **kwargs):
         """
         Context manager to facilitate writing the movie file.
 
-        ``*args, **kw`` are any parameters that should be passed to `setup`.
+        ``*args, **kwargs`` are any parameters that should be passed to `setup`.
         """
         if mpl.rcParams['savefig.bbox'] == 'tight':
             _log.info("Disabling savefig.bbox = 'tight', as it may cause "
@@ -939,7 +939,7 @@ def _stop(self, *args):
 
     def save(self, filename, writer=None, fps=None, dpi=None, codec=None,
              bitrate=None, extra_args=None, metadata=None, extra_anim=None,
-             savefig_kwargs=None, *, progress_callback=None):
+             savefig_kwargs=None, *, progress_callback=None, **kwargs):
         """
         Save the animation as a movie file by drawing every frame.
 
@@ -1006,6 +1006,11 @@ def func(current_frame: int, total_frames: int) -> Any
 
                 progress_callback = lambda i, n: print(f'Saving frame {i}/{n}')
 
+        **kwargs :
+            `AbstractMovieWriter` subclasses can specify additional parameters in their
+            `~.AbstractMovieWriter.setup` methods. Additional keyword arguments
+            are passed to these setup functions.
+
         Notes
         -----
         *fps*, *codec*, *bitrate*, *extra_args* and *metadata* are used to
@@ -1092,7 +1097,7 @@ def _pre_composite_to_white(color):
         # canvas._is_saving = True makes the draw_event animation-starting
         # callback a no-op; canvas.manager = None prevents resizing the GUI
         # widget (both are likewise done in savefig()).
-        with (writer.saving(self._fig, filename, dpi),
+        with (writer.saving(self._fig, filename, dpi, **kwargs),
               cbook._setattr_cm(self._fig.canvas, _is_saving=True, manager=None)):
             if not writer._supports_transparency():
                 facecolor = savefig_kwargs.get('facecolor',

lib/matplotlib/tests/test_animation.py

@@ -56,11 +56,12 @@ class NullMovieWriter(animation.AbstractMovieWriter):
     it cannot be added to the 'writers' registry.
     """
 
-    def setup(self, fig, outfile, dpi, *args):
+    def setup(self, fig, outfile, dpi, *args, **kwargs):
         self.fig = fig
         self.outfile = outfile
         self.dpi = dpi
         self.args = args
+        self.kwargs = kwargs
         self._count = 0
 
     def grab_frame(self, **savefig_kwargs):
@@ -79,15 +80,17 @@ def test_null_movie_writer(anim):
     filename = "unused.null"
     dpi = 50
     savefig_kwargs = dict(foo=0)
+    setup_kwargs = dict(bar=1)
     writer = NullMovieWriter()
 
     anim.save(filename, dpi=dpi, writer=writer,
-              savefig_kwargs=savefig_kwargs)
+              savefig_kwargs=savefig_kwargs, **setup_kwargs)
 
     assert writer.fig == plt.figure(1)  # The figure used by anim fixture
     assert writer.outfile == filename
     assert writer.dpi == dpi
     assert writer.args == ()
+    assert writer.kwargs == setup_kwargs
     # we enrich the savefig kwargs to ensure we composite transparent
     # output to an opaque background
     for k, v in savefig_kwargs.items():
@@ -187,8 +190,18 @@ def test_save_animation_smoketest(tmp_path, writer, frame_format, output, anim):
         dpi = 100.
         codec = 'h264'
 
+    setup_kwargs = {}
+    if writer.endswith('_file'):
+        assert len([*tmp_path.glob('example*')]) == 0
+        # Test that kwargs are passed to `MovieFileWriter.setup` so that the automatic
+        # temporary directory isn't used.
+        setup_kwargs['frame_prefix'] = str(tmp_path / 'example')
+
     anim.save(tmp_path / output, fps=30, writer=writer, bitrate=500, dpi=dpi,
-              codec=codec)
+              codec=codec, **setup_kwargs)
+
+    if writer.endswith('_file'):
+        assert len([*tmp_path.glob('example*')]) == anim._save_count
 
     del anim