Merge pull request #13564 from SBCV/animation_progress

timhoffm · web-flow · commit 4d29cf2a4185 · 2019-03-29T07:40:55.000+01:00
Add an option to log progress while saving animations
diff --git a/doc/users/next_whats_new/2019-03-18-SB.rst b/doc/users/next_whats_new/2019-03-18-SB.rst
@@ -0,0 +1,5 @@
+*progress_callback* argument to ani.save()
+-------------------------------------------------
+
+The method .FuncAnimation.save() gained an optional 
+*progress_callback* argument to notify the saving progress.
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -957,7 +957,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):
+             savefig_kwargs=None, *, progress_callback=None):
         """
         Save the animation as a movie file by drawing every frame.
 
@@ -1013,6 +1013,22 @@ class to use, such as 'ffmpeg'. If ``None``, defaults to
            on to the `savefig` command which is called repeatedly to
            save the individual frames.
 
+        progress_callback : function, optional
+            A callback function that will be called for every frame to notify
+            the saving progress. It must have the signature ::
+
+                def func(current_frame: int, total_frames: int) -> Any
+
+            where *current_frame* is the current frame number and
+            *total_frames* is the total number of frames to be saved.
+            *total_frames* is set to None, if the total number of frames can
+            not be determined. Return values may exist but are ignored.
+
+            Example code to write the progress to stdout::
+
+                progress_callback =\
+                    lambda i, n: print(f'Saving frame {i} of {n}')
+
         Notes
         -----
         *fps*, *codec*, *bitrate*, *extra_args* and *metadata* are used to
@@ -1111,10 +1127,19 @@ class to use, such as 'ffmpeg'. If ``None``, defaults to
                 for anim in all_anim:
                     # Clear the initial frame
                     anim._init_draw()
+                frame_number = 0
+                save_count_list = [a.save_count for a in all_anim]
+                if None in save_count_list:
+                    total_frames = None
+                else:
+                    total_frames = sum(save_count_list)
                 for data in zip(*[a.new_saved_frame_seq() for a in all_anim]):
                     for anim, d in zip(all_anim, data):
                         # TODO: See if turning off blit is really necessary
                         anim._draw_next_frame(d, blit=False)
+                        if progress_callback is not None:
+                            progress_callback(frame_number, total_frames)
+                            frame_number += 1
                     writer.grab_frame(**savefig_kwargs)
 
         # Reconnect signal for first draw if necessary
