Simplify _png extension by handling file open/close in Python.

Make functions in the (private) _png extension only take file-like
objects as arguments, rather than "filename or file-like".  Filenames
were anyways handled by calling the python C-API to open/close the
files, so we may as well just do this from Python where a simple
`with open(...):` does the job.  The public API is not changed.

Author: anntzer

lib/matplotlib/backends/backend_pgf.py

@@ -639,7 +639,8 @@ def draw_image(self, gc, x, y, im, transform=None):
         fname = os.path.splitext(os.path.basename(self.fh.name))[0]
         fname_img = "%s-img%d.png" % (fname, self.image_counter)
         self.image_counter += 1
-        _png.write_png(im[::-1], os.path.join(path, fname_img))
+        with pathlib.Path(path, fname_img).open("wb") as file:
+            _png.write_png(im[::-1], file)
 
         # reference the image in the pgf picture
         writeln(self.fh, r"\begin{pgfscope}")

lib/matplotlib/backends/backend_svg.py

@@ -821,20 +821,20 @@ def draw_image(self, gc, x, y, im, transform=None):
         if url is not None:
             self.writer.start('a', attrib={'xlink:href': url})
         if rcParams['svg.image_inline']:
-            bytesio = io.BytesIO()
-            _png.write_png(im, bytesio)
-            oid = oid or self._make_id('image', bytesio.getvalue())
+            buf = _png.write_png(im, None)
+            oid = oid or self._make_id('image', buf)
             attrib['xlink:href'] = (
                 "data:image/png;base64,\n" +
-                base64.b64encode(bytesio.getvalue()).decode('ascii'))
+                base64.b64encode(buf).decode('ascii'))
         else:
             if self.basename is None:
                 raise ValueError("Cannot save image data to filesystem when "
                                  "writing SVG to an in-memory buffer")
             filename = '{}.image{}.png'.format(
                 self.basename, next(self._image_counter))
             _log.info('Writing image file for inclusion: %s', filename)
-            _png.write_png(im, filename)
+            with open(filename, 'wb') as file:
+                _png.write_png(im, file)
             oid = oid or 'Im_' + self._make_id('image', filename)
             attrib['xlink:href'] = filename
 

lib/matplotlib/image.py

@@ -645,7 +645,8 @@ def write_png(self, fname):
         from matplotlib import _png
         im = self.to_rgba(self._A[::-1] if self.origin == 'lower' else self._A,
                           bytes=True, norm=True)
-        _png.write_png(im, fname)
+        with open(fname, "wb") as file:
+            _png.write_png(im, file)
 
     def set_data(self, A):
         """
@@ -1382,12 +1383,6 @@ def imread(fname, format=None):
 
     .. _Pillow documentation: http://pillow.readthedocs.io/en/latest/
     """
-
-    def read_png(*args, **kwargs):
-        from matplotlib import _png
-        return _png.read_png(*args, **kwargs)
-
-    handlers = {'png': read_png, }
     if format is None:
         if isinstance(fname, str):
             parsed = urllib.parse.urlparse(fname)
@@ -1412,34 +1407,24 @@ def read_png(*args, **kwargs):
             ext = 'png'
     else:
         ext = format
-
-    if ext not in handlers:  # Try to load the image with PIL.
-        try:
+    if ext != 'png':
+        try:  # Try to load the image with PIL.
             from PIL import Image
         except ImportError:
-            raise ValueError('Only know how to handle extensions: %s; '
-                             'with Pillow installed matplotlib can handle '
-                             'more images' % list(handlers))
+            raise ValueError('Only know how to handle PNG; with Pillow '
+                             'installed, Matplotlib can handle more images')
         with Image.open(fname) as image:
             return pil_to_array(image)
-
-    handler = handlers[ext]
-
-    # To handle Unicode filenames, we pass a file object to the PNG
-    # reader extension, since Python handles them quite well, but it's
-    # tricky in C.
+    from matplotlib import _png
     if isinstance(fname, str):
         parsed = urllib.parse.urlparse(fname)
         # If fname is a URL, download the data
         if len(parsed.scheme) > 1:
             from urllib import request
             fd = BytesIO(request.urlopen(fname).read())
-            return handler(fd)
-        else:
-            with open(fname, 'rb') as fd:
-                return handler(fd)
-    else:
-        return handler(fname)
+            return _png.read_png(fd)
+    with cbook.open_file_cm(fname, "rb") as file:
+        return _png.read_png(file)
 
 
 def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
@@ -1516,7 +1501,8 @@ def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
             arr = arr[::-1]
         rgba = sm.to_rgba(arr, bytes=True)
         if format == "png" and pil_kwargs is None:
-            _png.write_png(rgba, fname, dpi=dpi, metadata=metadata)
+            with cbook.open_file_cm(fname, "wb") as file:
+                _png.write_png(rgba, file, dpi=dpi, metadata=metadata)
         else:
             try:
                 from PIL import Image

lib/matplotlib/mathtext.py

@@ -3433,7 +3433,8 @@ def to_png(self, filename, texstr, color='black', dpi=120, fontsize=14):
         from matplotlib import _png
         rgba, depth = self.to_rgba(
             texstr, color=color, dpi=dpi, fontsize=fontsize)
-        _png.write_png(rgba, filename)
+        with open(filename, "wb") as file:
+            _png.write_png(rgba, file)
         return depth
 
     def get_depth(self, texstr, dpi=120, fontsize=14):

lib/matplotlib/testing/compare.py

@@ -384,10 +384,10 @@ def compare_images(expected, actual, tol, in_decorator=False):
         expected = convert(expected, True)
 
     # open the image files and remove the alpha channel (if it exists)
-    expected_image = _png.read_png_int(expected)
-    actual_image = _png.read_png_int(actual)
-    expected_image = expected_image[:, :, :3]
-    actual_image = actual_image[:, :, :3]
+    with open(expected, "rb") as expected_file:
+        expected_image = _png.read_png_int(expected_file)[:, :, :3]
+    with open(actual, "rb") as actual_file:
+        actual_image = _png.read_png_int(actual_file)[:, :, :3]
 
     actual_image, expected_image = crop_to_same(
         actual, actual_image, expected, expected_image)
@@ -438,8 +438,10 @@ def save_diff_image(expected, actual, output):
     '''
     # Drop alpha channels, similarly to compare_images.
     from matplotlib import _png
-    expected_image = _png.read_png(expected)[..., :3]
-    actual_image = _png.read_png(actual)[..., :3]
+    with open(expected, "rb") as expected_file:
+        expected_image = _png.read_png(expected_file)[..., :3]
+    with open(actual, "rb") as actual_file:
+        actual_image = _png.read_png(actual_file)[..., :3]
     actual_image, expected_image = crop_to_same(
         actual, actual_image, expected, expected_image)
     expected_image = np.array(expected_image).astype(float)
@@ -465,4 +467,5 @@ def save_diff_image(expected, actual, output):
     # Hard-code the alpha channel to fully solid
     save_image_np[:, :, 3] = 255
 
-    _png.write_png(save_image_np, output)
+    with open(output, "wb") as output_file:
+        _png.write_png(save_image_np, output_file)

lib/matplotlib/tests/test_png.py

@@ -1,6 +1,8 @@
 from io import BytesIO
 import glob
 import os
+from pathlib import Path
+
 import numpy as np
 import pytest
 
@@ -33,9 +35,9 @@ def test_pngsuite():
 
 def test_imread_png_uint16():
     from matplotlib import _png
-    img = _png.read_png_int(os.path.join(os.path.dirname(__file__),
-                                     'baseline_images/test_png/uint16.png'))
-
+    with (Path(__file__).parent
+          / 'baseline_images/test_png/uint16.png').open('rb') as file:
+        img = _png.read_png_int(file)
     assert (img.dtype == np.uint16)
     assert np.sum(img.flatten()) == 134184960
 

lib/matplotlib/texmanager.py

@@ -403,7 +403,8 @@ def get_grey(self, tex, fontsize=None, dpi=None):
         alpha = self.grey_arrayd.get(key)
         if alpha is None:
             pngfile = self.make_png(tex, fontsize, dpi)
-            X = _png.read_png(os.path.join(self.texcache, pngfile))
+            with open(os.path.join(self.texcache, pngfile), "rb") as file:
+                X = _png.read_png(file)
             self.grey_arrayd[key] = alpha = X[:, :, -1]
         return alpha