Allow passing arguments to PIL.Image.save().

anntzer · anntzer · commit 71eca0255fe3 · 2019-01-04T10:28:52.000+01:00
diff --git a/doc/users/next_whats_new/2018-01-03-AL.rst b/doc/users/next_whats_new/2018-01-03-AL.rst
@@ -0,0 +1,6 @@
+*pil_kwargs* argument to savefig
+````````````````````````````````
+
+Matplotlib uses Pillow to handle saving to the JPEG and TIFF formats.  The
+`~Figure.savefig()` function gained a *pil_kwargs* keyword argument, which can
+be used to forward arguments to Pillow's ``Image.save()``.
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -450,6 +450,7 @@ def print_raw(self, filename_or_obj, *args, **kwargs):
         with cbook._setattr_cm(renderer, dpi=self.figure.dpi), \
                 cbook.open_file_cm(filename_or_obj, "wb") as fh:
             fh.write(renderer._renderer.buffer_rgba())
+
     print_rgba = print_raw
 
     def print_png(self, filename_or_obj, *args, **kwargs):
@@ -507,7 +508,7 @@ def print_png(self, filename_or_obj, *args, **kwargs):
         with cbook._setattr_cm(renderer, dpi=self.figure.dpi), \
                 cbook.open_file_cm(filename_or_obj, "wb") as fh:
             _png.write_png(renderer._renderer, fh,
-                            self.figure.dpi, metadata=metadata)
+                           self.figure.dpi, metadata=metadata)
 
     def print_to_buffer(self):
         FigureCanvasAgg.draw(self)
@@ -517,8 +518,13 @@ def print_to_buffer(self):
                     (int(renderer.width), int(renderer.height)))
 
     if _has_pil:
-        # add JPEG support
-        def print_jpg(self, filename_or_obj, *args, dryrun=False, **kwargs):
+
+        # Note that these methods should typically be called via savefig() and
+        # print_figure(), and the latter ensures that `self.figure.dpi` already
+        # matches the dpi kwarg (if any).
+
+        def print_jpg(self, filename_or_obj, *args, dryrun=False,
+                      pil_kwargs=None, **kwargs):
             """
             Write the figure to a JPEG file.
 
@@ -543,6 +549,11 @@ def print_jpg(self, filename_or_obj, *args, dryrun=False, **kwargs):
             progressive : bool
                 If present, indicates that this image
                 should be stored as a progressive JPEG file.
+
+            pil_kwargs : dict, optional
+                Additional keyword arguments that are passed to Pillow when
+                saving the figure.  These take precedence over *quality*,
+                *optimize* and *progressive*.
             """
             buf, size = self.print_to_buffer()
             if dryrun:
@@ -554,25 +565,29 @@ def print_jpg(self, filename_or_obj, *args, dryrun=False, **kwargs):
             color = tuple([int(x * 255) for x in rgba[:3]])
             background = Image.new('RGB', size, color)
             background.paste(image, image)
-            options = {k: kwargs[k]
-                       for k in ['quality', 'optimize', 'progressive', 'dpi']
-                       if k in kwargs}
-            options.setdefault('quality', rcParams['savefig.jpeg_quality'])
-            if 'dpi' in options:
-                # Set the same dpi in both x and y directions
-                options['dpi'] = (options['dpi'], options['dpi'])
-
-            return background.save(filename_or_obj, format='jpeg', **options)
+            if pil_kwargs is None:
+                pil_kwargs = {}
+            for k in ["quality", "optimize", "progressive"]:
+                if k in kwargs:
+                    pil_kwargs.setdefault(k, kwargs[k])
+            pil_kwargs.setdefault("quality", rcParams["savefig.jpeg_quality"])
+            pil_kwargs.setdefault("dpi", (self.figure.dpi, self.figure.dpi))
+            return background.save(
+                filename_or_obj, format='jpeg', **pil_kwargs)
+
         print_jpeg = print_jpg
 
-        # add TIFF support
-        def print_tif(self, filename_or_obj, *args, dryrun=False, **kwargs):
+        def print_tif(self, filename_or_obj, *args, dryrun=False,
+                      pil_kwargs=None, **kwargs):
             buf, size = self.print_to_buffer()
             if dryrun:
                 return
             image = Image.frombuffer('RGBA', size, buf, 'raw', 'RGBA', 0, 1)
-            dpi = (self.figure.dpi, self.figure.dpi)
-            return image.save(filename_or_obj, format='tiff', dpi=dpi)
+            if pil_kwargs is None:
+                pil_kwargs = {}
+            pil_kwargs.setdefault("dpi", (self.figure.dpi, self.figure.dpi))
+            return image.save(filename_or_obj, format='tiff', **pil_kwargs)
+
         print_tiff = print_tif
 
 
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -2103,7 +2103,12 @@ def savefig(self, fname, *, frameon=None, transparent=None, **kwargs):
               `~.backend_pdf.PdfPages`.
             - 'eps' and 'ps' with PS backend: Only 'Creator' is supported.
 
+        pil_kwargs : dict, optional
+            Additional keyword arguments that are passed to Pillow when saving
+            the figure.  Only applicable for formats that are saved using
+            Pillow, i.e. JPEG and TIFF.
         """
+
         kwargs.setdefault('dpi', rcParams['savefig.dpi'])
         if frameon is None:
             frameon = rcParams['savefig.frameon']
diff --git a/lib/matplotlib/tests/test_agg.py b/lib/matplotlib/tests/test_agg.py
@@ -239,3 +239,14 @@ def test_jpeg_dpi():
     plt.savefig(buf, format="jpg", dpi=200)
     im = Image.open(buf)
     assert im.info['dpi'] == (200, 200)
+
+
+def test_pil_kwargs():
+    Image = pytest.importorskip("PIL.Image")
+    from PIL.TiffTags import TAGS_V2 as TAGS
+    buf = io.BytesIO()
+    pil_kwargs = {"description": "test image"}
+    plt.figure().savefig(buf, format="tiff", pil_kwargs=pil_kwargs)
+    im = Image.open(buf)
+    tags = {TAGS[k].name: v for k, v in im.tag_v2.items()}
+    assert tags["ImageDescription"] == "test image"
