Merge pull request #1442 from maxalbert/add_savefig_kwargs_to_animation_save_method

dopplershift · dopplershift · commit 8b1395223315 · 2012-11-16T10:16:46.000-08:00
Add savefig_kwargs to Animation.save() method
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -184,17 +184,19 @@ def finish(self):
         'Finish any processing for writing the movie.'
         self.cleanup()
 
-    def grab_frame(self):
+    def grab_frame(self, **savefig_kwargs):
         '''
         Grab the image information from the figure and save as a movie frame.
+        All keyword arguments in savefig_kwargs are passed on to the 'savefig'
+        command that saves the figure.
         '''
         verbose.report('MovieWriter.grab_frame: Grabbing frame.',
                        level='debug')
         try:
             # Tell the figure to save its data to the sink, using the
             # frame format and dpi.
             self.fig.savefig(self._frame_sink(), format=self.frame_format,
-                dpi=self.dpi)
+                dpi=self.dpi, **savefig_kwargs)
         except RuntimeError:
             out, err = self._proc.communicate()
             verbose.report('MovieWriter -- Error running proc:\n%s\n%s' % (out,
@@ -545,7 +547,8 @@ def _stop(self, *args):
         self.event_source = None
 
     def save(self, filename, writer=None, fps=None, dpi=None, codec=None,
-             bitrate=None, extra_args=None, metadata=None, extra_anim=None):
+             bitrate=None, extra_args=None, metadata=None, extra_anim=None,
+             savefig_kwargs=None):
         '''
         Saves a movie file by drawing every frame.
 
@@ -586,7 +589,32 @@ def save(self, filename, writer=None, fps=None, dpi=None, codec=None,
         `matplotlib.Figure` instance. Also, animation frames will just be
         simply combined, so there should be a 1:1 correspondence between
         the frames from the different animations.
+
+        *savefig_kwargs* is a dictionary containing keyword arguments to be
+        passed on to the 'savefig' command which is called repeatedly to save
+        the individual frames. This can be used to set tight bounding boxes,
+        for example.
         '''
+        if savefig_kwargs is None:
+            savefig_kwargs = {}
+
+        # FIXME: Using 'bbox_inches' doesn't currently work with writers that pipe
+        # the data to the command because this requires a fixed frame size (see
+        # Ryan May's reply in this thread: [1]). Thus we drop the 'bbox_inches'
+        # argument if it exists in savefig_kwargs.
+        #
+        # [1] http://matplotlib.1069221.n5.nabble.com/Animation-class-let-save-accept-kwargs-which-are-passed-on-to-savefig-td39627.html
+        #
+        if savefig_kwargs.has_key('bbox_inches'):
+            if not (writer in ['ffmpeg_file', 'mencoder_file'] or
+                    isinstance(writer, (FFMpegFileWriter, MencoderFileWriter))):
+                print("Warning: discarding the 'bbox_inches' argument in " \
+                          "'savefig_kwargs' as it is only currently supported " \
+                          "with the writers 'ffmpeg_file' and 'mencoder_file' " \
+                          "(writer used: '{}').".format(writer if isinstance(writer, str)
+                                                        else writer.__class__.__name__))
+                savefig_kwargs.pop('bbox_inches')
+
         # Need to disconnect the first draw callback, since we'll be doing
         # draws. Otherwise, we'll end up starting the animation.
         if self._first_draw_id is not None:
@@ -645,7 +673,7 @@ def save(self, filename, writer=None, fps=None, dpi=None, codec=None,
                 for anim, d in zip(all_anim, data):
                     #TODO: Need to see if turning off blit is really necessary
                     anim._draw_next_frame(d, blit=False)
-                writer.grab_frame()
+                writer.grab_frame(**savefig_kwargs)
 
         # Reconnect signal for first draw if necessary
         if reconnect_first_draw:
