From 9c7f926d6d70a994c3710ce570bb9807519ef090 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Sun, 6 Jan 2019 13:37:46 +0100 Subject: [PATCH] Allow saving PNGs through Pillow as well. --- doc/users/next_whats_new/2018-01-03-AL.rst | 4 ++ lib/matplotlib/backends/backend_agg.py | 52 ++++++++++++++++------ lib/matplotlib/figure.py | 3 +- lib/matplotlib/tests/test_agg.py | 13 +++++- 4 files changed, 57 insertions(+), 15 deletions(-) diff --git a/doc/users/next_whats_new/2018-01-03-AL.rst b/doc/users/next_whats_new/2018-01-03-AL.rst index 7f3da98c8a7b..de6787834c09 100644 --- 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 index 5fce6803c50a..ec917da264b6 100644 --- 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 index 136ba3999499..5c30805a7d73 100644 --- 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 index 13e46aeeccb5..339ab2118125 100644 --- 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()