Factor common parts of saving to different formats using pillow.

anntzer · anntzer · commit 3a87bb7bfd1e · 2021-10-18T12:29:44.000+02:00
We don't need to drop the alpha channel for jpeg anymore as imsave
handles that.

In imsave, also clarify that `sm` is not going to be used in the
already-usable-rgba case.
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -487,6 +487,16 @@ def print_raw(self, filename_or_obj, *args):
 
     print_rgba = print_raw
 
+    def _print_pil(self, filename_or_obj, fmt, pil_kwargs, metadata=None):
+        """
+        Draw the canvas, then save it using `.image.imsave` (to which
+        *pil_kwargs* and *metadata* are forwarded).
+        """
+        FigureCanvasAgg.draw(self)
+        mpl.image.imsave(
+            filename_or_obj, self.buffer_rgba(), format=fmt, origin="upper",
+            dpi=self.figure.dpi, metadata=metadata, pil_kwargs=pil_kwargs)
+
     @_check_savefig_extra_args
     @_api.delete_parameter("3.5", "args")
     def print_png(self, filename_or_obj, *args,
@@ -537,10 +547,7 @@ def print_png(self, filename_or_obj, *args,
             If the 'pnginfo' key is present, it completely overrides
             *metadata*, including the default 'Software' key.
         """
-        FigureCanvasAgg.draw(self)
-        mpl.image.imsave(
-            filename_or_obj, self.buffer_rgba(), format="png", origin="upper",
-            dpi=self.figure.dpi, metadata=metadata, pil_kwargs=pil_kwargs)
+        self._print_pil(filename_or_obj, "png", pil_kwargs, metadata)
 
     def print_to_buffer(self):
         FigureCanvasAgg.draw(self)
@@ -573,26 +580,15 @@ def print_jpg(self, filename_or_obj, *args, pil_kwargs=None, **kwargs):
         r, g, b, a = mcolors.to_rgba(self.figure.get_facecolor())
         try:
             self.figure.set_facecolor(a * np.array([r, g, b]) + 1 - a)
-            FigureCanvasAgg.draw(self)
+            self._print_pil(filename_or_obj, "jpeg", pil_kwargs)
         finally:
             self.figure.set_facecolor((r, g, b, a))
-        if pil_kwargs is None:
-            pil_kwargs = {}
-        pil_kwargs.setdefault("dpi", (self.figure.dpi, self.figure.dpi))
-        # Drop alpha channel now.
-        return (Image.fromarray(np.asarray(self.buffer_rgba())[..., :3])
-                .save(filename_or_obj, format='jpeg', **pil_kwargs))
 
     print_jpeg = print_jpg
 
     @_check_savefig_extra_args
     def print_tif(self, filename_or_obj, *, pil_kwargs=None):
-        FigureCanvasAgg.draw(self)
-        if pil_kwargs is None:
-            pil_kwargs = {}
-        pil_kwargs.setdefault("dpi", (self.figure.dpi, self.figure.dpi))
-        return (Image.fromarray(np.asarray(self.buffer_rgba()))
-                .save(filename_or_obj, format='tiff', **pil_kwargs))
+        self._print_pil(filename_or_obj, "tiff", pil_kwargs)
 
     print_tiff = print_tif
 
@@ -612,11 +608,7 @@ def print_webp(self, filename_or_obj, *, pil_kwargs=None):
             Additional keyword arguments that are passed to
             `PIL.Image.Image.save` when saving the figure.
         """
-        FigureCanvasAgg.draw(self)
-        if pil_kwargs is None:
-            pil_kwargs = {}
-        return (Image.fromarray(np.asarray(self.buffer_rgba()))
-                .save(filename_or_obj, format='webp', **pil_kwargs))
+        self._print_pil(filename_or_obj, "webp", pil_kwargs)
 
 
 @_Backend.export
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -1601,20 +1601,20 @@ def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
     else:
         # Don't bother creating an image; this avoids rounding errors on the
         # size when dividing and then multiplying by dpi.
-        sm = cm.ScalarMappable(cmap=cmap)
-        sm.set_clim(vmin, vmax)
         if origin is None:
             origin = mpl.rcParams["image.origin"]
         if origin == "lower":
             arr = arr[::-1]
         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
+            # Such an ``arr`` would also be handled fine by sm.to_rgba below
+            # (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:
+            sm = cm.ScalarMappable(cmap=cmap)
+            sm.set_clim(vmin, vmax)
             rgba = sm.to_rgba(arr, bytes=True)
         if pil_kwargs is None:
             pil_kwargs = {}
