Pillow animation writer.

anntzer · anntzer · commit e0e346d7ce78 · 2018-01-22T09:54:19.000-08:00
diff --git a/doc/api/animation_api.rst b/doc/api/animation_api.rst
@@ -153,10 +153,20 @@ Examples
 Writer Classes
 ==============
 
-The provided writers fall into two broad categories: pipe-based and
-file-based.  The pipe-based writers stream the captured frames over a
-pipe to an external process.  The pipe-based variants tend to be more
-performant, but may not work on all systems.
+The provided writers fall into a few broad categories.
+
+The Pillow writer relies on the Pillow library to write the animation, keeping
+all data in memory.
+
+.. autosummary::
+   :toctree: _as_gen
+   :nosignatures:
+
+   PillowWriter
+
+The pipe-based writers stream the captured frames over a pipe to an external
+process.  The pipe-based variants tend to be more performant, but may not work
+on all systems.
 
 .. autosummary::
    :toctree: _as_gen
@@ -166,9 +176,9 @@ performant, but may not work on all systems.
    ImageMagickFileWriter
    AVConvWriter
 
-Alternatively the file-based writers save temporary files for each
-frame which are stitched into a single file at the end.  Although
-slower, these writers can be easier to debug.
+The file-based writers save temporary files for each frame which are stitched
+into a single file at the end.  Although slower, these writers can be easier to
+debug.
 
 .. autosummary::
    :toctree: _as_gen
diff --git a/doc/users/next_whats_new/180112-AL-pillowanim.rst b/doc/users/next_whats_new/180112-AL-pillowanim.rst
@@ -0,0 +1,20 @@
+Writing animations with Pillow
+``````````````````````````````
+It is now possible to use Pillow as an animation writer.  Supported output
+formats are currently gif (Pillow>=3.4) and webp (Pillow>=5.0).  Use e.g. as ::
+
+   from __future__ import division
+
+   from matplotlib import pyplot as plt
+   from matplotlib.animation import FuncAnimation, PillowWriter
+
+   fig, ax = plt.subplots()
+   line, = plt.plot([0, 1])
+
+   def animate(i):
+      line.set_ydata([0, i / 20])
+      return [line]
+
+   anim = FuncAnimation(fig, animate, 20, blit=True)
+   anim.save("movie.gif", writer=PillowWriter(fps=24))
+   plt.show()
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -25,6 +25,7 @@
 
 import abc
 import contextlib
+from io import BytesIO
 import itertools
 import logging
 import os
@@ -43,10 +44,8 @@
 
 if six.PY2:
     from base64 import encodestring as encodebytes
-    from cStringIO import StringIO as BytesIO
 else:
     from base64 import encodebytes
-    from io import BytesIO
 
 
 _log = logging.getLogger(__name__)
@@ -575,6 +574,45 @@ def cleanup(self):
                 os.remove(fname)
 
 
+@writers.register('pillow')
+class PillowWriter(MovieWriter):
+    @classmethod
+    def isAvailable(cls):
+        try:
+            import PIL
+        except ImportError:
+            return False
+        return True
+
+    def __init__(self, *args, **kwargs):
+        if kwargs.get("extra_args") is None:
+            kwargs["extra_args"] = ()
+        super(PillowWriter, self).__init__(*args, **kwargs)
+
+    def setup(self, fig, outfile, dpi=None):
+        self._frames = []
+        self._outfile = outfile
+        self._dpi = dpi
+        self._fig = fig
+
+    def grab_frame(self, **savefig_kwargs):
+        from PIL import Image
+        buf = BytesIO()
+        self._fig.savefig(buf, **dict(savefig_kwargs, format="rgba"))
+        renderer = self._fig.canvas.get_renderer()
+        # Using frombuffer / getbuffer may be slightly more efficient, but
+        # Py3-only.
+        self._frames.append(Image.frombytes(
+            "RGBA",
+            (int(renderer.width), int(renderer.height)),
+            buf.getvalue()))
+
+    def finish(self):
+        self._frames[0].save(
+            self._outfile, save_all=True, append_images=self._frames[1:],
+            duration=int(1000 / self.fps))
+
+
 # Base class of ffmpeg information. Has the config keys and the common set
 # of arguments that controls the *output* side of things.
 class FFMpegBase(object):
@@ -1138,8 +1176,8 @@ class to use, such as 'ffmpeg'. If ``None``, defaults to
 
         if 'bbox_inches' in savefig_kwargs:
             _log.warning("Warning: discarding the 'bbox_inches' argument in "
-                          "'savefig_kwargs' as it may cause frame size "
-                          "to vary, which is inappropriate for animation.")
+                         "'savefig_kwargs' as it may cause frame size "
+                         "to vary, which is inappropriate for animation.")
             savefig_kwargs.pop('bbox_inches')
 
         # Create a new sequence of frames for saved data. This is different
@@ -1150,16 +1188,15 @@ class to use, such as 'ffmpeg'. If ``None``, defaults to
         # allow for this non-existent use case or find a way to make it work.
         with rc_context():
             if rcParams['savefig.bbox'] == 'tight':
-                _log.info("Disabling savefig.bbox = 'tight', as it "
-                               "may cause frame size to vary, which "
-                               "is inappropriate for animation.")
+                _log.info("Disabling savefig.bbox = 'tight', as it may cause "
+                          "frame size to vary, which is inappropriate for "
+                          "animation.")
                 rcParams['savefig.bbox'] = None
             with writer.saving(self._fig, filename, dpi):
                 for anim in all_anim:
                     # Clear the initial frame
                     anim._init_draw()
-                for data in zip(*[a.new_saved_frame_seq()
-                                  for a in all_anim]):
+                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)
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -117,6 +117,7 @@ def isAvailable(self):
     ('avconv_file', 'movie.mp4'),
     ('imagemagick', 'movie.gif'),
     ('imagemagick_file', 'movie.gif'),
+    ('pillow', 'movie.gif'),
     ('html', 'movie.html'),
     ('null', 'movie.null')
 ]

Original file line number	Diff line number	Diff line change
`@@ -117,6 +117,7 @@ def isAvailable(self):`
`117`	`117`	`('avconv_file', 'movie.mp4'),`
`118`	`118`	`('imagemagick', 'movie.gif'),`
`119`	`119`	`('imagemagick_file', 'movie.gif'),`
	`120`	`+ ('pillow', 'movie.gif'),`
`120`	`121`	`('html', 'movie.html'),`
`121`	`122`	`('null', 'movie.null')`
`122`	`123`	`]`