Allow saving PNGs through Pillow as well.

anntzer · anntzer · commit 9c7f926d6d70 · 2019-01-17T19:04:07.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
@@ -4,3 +4,7 @@
 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 `PIL.Image.save()`.
+
+The *pil_kwargs* argument can also be used when saving to PNG.  In that case,
+Matplotlib also uses Pillow's `PIL.Image.save()` instead of going through its
+own builtin PNG support.
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -453,7 +453,9 @@ def print_raw(self, filename_or_obj, *args, **kwargs):
 
     print_rgba = print_raw
 
-    def print_png(self, filename_or_obj, *args, **kwargs):
+    def print_png(self, filename_or_obj, *args,
+                  metadata=None, pil_kwargs=None,
+                  **kwargs):
         """
         Write the figure to a PNG file.
 
@@ -494,21 +496,45 @@ def print_png(self, filename_or_obj, *args, **kwargs):
             .. _PNG specification: \
                 https://www.w3.org/TR/2003/REC-PNG-20031110/#11keywords
 
+        pil_kwargs : dict, optional
+            If set to a non-None value, use Pillow to save the figure instead
+            of Matplotlib's builtin PNG support, and pass these keyword
+            arguments to `PIL.Image.save`.
+
+            If the 'pnginfo' key is present, it completely overrides
+            *metadata*, including the default 'Software' key.
         """
-        FigureCanvasAgg.draw(self)
-        renderer = self.get_renderer()
 
-        version_str = (
-            'matplotlib version ' + __version__ + ', http://matplotlib.org/')
-        metadata = OrderedDict({'Software': version_str})
-        user_metadata = kwargs.pop("metadata", None)
-        if user_metadata is not None:
-            metadata.update(user_metadata)
+        if metadata is None:
+            metadata = {}
+        metadata = {
+            "Software":
+                f"matplotlib version{__version__}, http://matplotlib.org/",
+            **metadata,
+        }
 
-        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)
+        if pil_kwargs is not None:
+            from PIL import Image
+            from PIL.PngImagePlugin import PngInfo
+            buf, size = self.print_to_buffer()
+            # Only use the metadata kwarg if pnginfo is not set, because the
+            # semantics of duplicate keys in pnginfo is unclear.
+            if "pnginfo" not in pil_kwargs:
+                pnginfo = PngInfo()
+                for k, v in metadata.items():
+                    pnginfo.add_text(k, v)
+                pil_kwargs["pnginfo"] = pnginfo
+            pil_kwargs.setdefault("dpi", (self.figure.dpi, self.figure.dpi))
+            (Image.frombuffer("RGBA", size, buf, "raw", "RGBA", 0, 1)
+             .save(filename_or_obj, format="png", **pil_kwargs))
+
+        else:
+            FigureCanvasAgg.draw(self)
+            renderer = self.get_renderer()
+            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)
 
     def print_to_buffer(self):
         FigureCanvasAgg.draw(self)
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -2106,7 +2106,8 @@ def savefig(self, fname, *, frameon=None, transparent=None, **kwargs):
         pil_kwargs : dict, optional
             Additional keyword arguments that are passed to `PIL.Image.save`
             when saving the figure.  Only applicable for formats that are saved
-            using Pillow, i.e. JPEG and TIFF.
+            using Pillow, i.e. JPEG, TIFF, and (if the keyword is set to a
+            non-None value) PNG.
         """
 
         kwargs.setdefault('dpi', rcParams['savefig.dpi'])
diff --git a/lib/matplotlib/tests/test_agg.py b/lib/matplotlib/tests/test_agg.py
@@ -241,7 +241,18 @@ def test_jpeg_dpi():
     assert im.info['dpi'] == (200, 200)
 
 
-def test_pil_kwargs():
+def test_pil_kwargs_png():
+    Image = pytest.importorskip("PIL.Image")
+    from PIL.PngImagePlugin import PngInfo
+    buf = io.BytesIO()
+    pnginfo = PngInfo()
+    pnginfo.add_text("Software", "test")
+    plt.figure().savefig(buf, format="png", pil_kwargs={"pnginfo": pnginfo})
+    im = Image.open(buf)
+    assert im.info["Software"] == "test"
+
+
+def test_pil_kwargs_tiff():
     Image = pytest.importorskip("PIL.Image")
     from PIL.TiffTags import TAGS_V2 as TAGS
     buf = io.BytesIO()
