Merge pull request #15435 from anntzer/reuse-imsave

QuLogic · web-flow · commit da3baa11a105 · 2020-05-29T17:52:27.000-04:00
Reuse png metadata handling of imsave() in FigureCanvasAgg.print_png().
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -33,7 +33,6 @@
 
 import numpy as np
 from PIL import Image
-from PIL.PngImagePlugin import PngInfo
 
 import matplotlib as mpl
 from matplotlib import cbook
@@ -491,8 +490,8 @@ def print_png(self, filename_or_obj, *args,
 
             Other keywords may be invented for other purposes.
 
-            If 'Software' is not given, an autogenerated value for matplotlib
-            will be used.
+            If 'Software' is not given, an autogenerated value for Matplotlib
+            will be used.  This can be removed by setting it to *None*.
 
             For more details see the `PNG specification`_.
 
@@ -505,27 +504,10 @@ def print_png(self, filename_or_obj, *args,
             If the 'pnginfo' key is present, it completely overrides
             *metadata*, including the default 'Software' key.
         """
-
-        if metadata is None:
-            metadata = {}
-        if pil_kwargs is None:
-            pil_kwargs = {}
-        metadata = {
-            "Software":
-                f"matplotlib version{mpl.__version__}, http://matplotlib.org/",
-            **metadata,
-        }
         FigureCanvasAgg.draw(self)
-        # 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.fromarray(np.asarray(self.buffer_rgba()))
-         .save(filename_or_obj, format="png", **pil_kwargs))
+        mpl.image.imsave(
+            filename_or_obj, self.buffer_rgba(), format="png", origin="upper",
+            dpi=self.figure.dpi, metadata=metadata, pil_kwargs=pil_kwargs)
 
     def print_to_buffer(self):
         FigureCanvasAgg.draw(self)
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -1515,18 +1515,37 @@ def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
             origin = mpl.rcParams["image.origin"]
         if origin == "lower":
             arr = arr[::-1]
-        rgba = sm.to_rgba(arr, bytes=True)
+        if (isinstance(arr, memoryview) and arr.format == "B"
+                and arr.ndim == 3 and arr.shape[-1] == 4):
+            # Such an ``arr`` would also be handled fine by sm.to_rgba (after
+            # casting with asarray), but it is useful to special-case it
+            # because that's what backend_agg passes, and can be in fact used
+            # as is, saving a few operations.
+            rgba = arr
+        else:
+            rgba = sm.to_rgba(arr, bytes=True)
         if pil_kwargs is None:
             pil_kwargs = {}
         pil_shape = (rgba.shape[1], rgba.shape[0])
         image = PIL.Image.frombuffer(
             "RGBA", pil_shape, rgba, "raw", "RGBA", 0, 1)
-        if format == "png" and metadata is not None:
-            # cf. backend_agg's print_png.
-            pnginfo = PIL.PngImagePlugin.PngInfo()
-            for k, v in metadata.items():
-                pnginfo.add_text(k, v)
-            pil_kwargs["pnginfo"] = pnginfo
+        if format == "png":
+            # Only use the metadata kwarg if pnginfo is not set, because the
+            # semantics of duplicate keys in pnginfo is unclear.
+            if "pnginfo" in pil_kwargs:
+                if metadata:
+                    cbook._warn_external("'metadata' is overridden by the "
+                                         "'pnginfo' entry in 'pil_kwargs'.")
+            else:
+                metadata = {
+                    "Software": (f"Matplotlib version{mpl.__version__}, "
+                                 f"https://matplotlib.org/"),
+                    **(metadata if metadata is not None else {}),
+                }
+                pil_kwargs["pnginfo"] = pnginfo = PIL.PngImagePlugin.PngInfo()
+                for k, v in metadata.items():
+                    if v is not None:
+                        pnginfo.add_text(k, v)
         if format in ["jpg", "jpeg"]:
             format = "jpeg"  # Pillow doesn't recognize "jpg".
             facecolor = mpl.rcParams["savefig.facecolor"]
