Merge pull request #10231 from anntzer/pathlike

tacaswell · web-flow · commit af1197d2d8ad · 2018-01-21T16:12:03.000-05:00
ENH: Support PathLike inputs.
diff --git a/doc/users/next_whats_new/2018-01-10-AL-pathlike.rst b/doc/users/next_whats_new/2018-01-10-AL-pathlike.rst
@@ -0,0 +1,6 @@
+PathLike support
+````````````````
+
+On Python 3.6+, `~matplotlib.pyplot.savefig`, `~matplotlib.pyplot.imsave`,
+`~matplotlib.pyplot.imread`, and animation writers now accept `os.PathLike`\s
+as input.
diff --git a/lib/matplotlib/animation.py b/lib/matplotlib/animation.py
@@ -358,8 +358,7 @@ def _run(self):
         # from a few configuration options.
         command = self._args()
         output = subprocess.PIPE
-        _log.info('MovieWriter.run: running command: %s',
-                       ' '.join(command))
+        _log.info('MovieWriter.run: running command: %s', command)
         self._proc = subprocess.Popen(command, shell=False,
                                       stdout=output, stderr=output,
                                       stdin=subprocess.PIPE,
@@ -847,12 +846,13 @@ def __init__(self, fps=30, codec=None, bitrate=None, extra_args=None,
                                          extra_args, metadata)
 
     def setup(self, fig, outfile, dpi, frame_dir=None):
-        if os.path.splitext(outfile)[-1] not in ['.html', '.htm']:
+        root, ext = os.path.splitext(outfile)
+        if ext not in ['.html', '.htm']:
             raise ValueError("outfile must be *.htm or *.html")
 
         if not self.embed_frames:
             if frame_dir is None:
-                frame_dir = outfile.rstrip('.html') + '_frames'
+                frame_dir = root + '_frames'
             if not os.path.exists(frame_dir):
                 os.makedirs(frame_dir)
             frame_prefix = os.path.join(frame_dir, 'frame')
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -510,11 +510,6 @@ def print_png(self, filename_or_obj, *args, **kwargs):
         renderer = self.get_renderer()
         original_dpi = renderer.dpi
         renderer.dpi = self.figure.dpi
-        if isinstance(filename_or_obj, six.string_types):
-            filename_or_obj = open(filename_or_obj, 'wb')
-            close = True
-        else:
-            close = False
 
         version_str = 'matplotlib version ' + __version__ + \
             ', http://matplotlib.org/'
@@ -524,11 +519,10 @@ def print_png(self, filename_or_obj, *args, **kwargs):
             metadata.update(user_metadata)
 
         try:
-            _png.write_png(renderer._renderer, filename_or_obj,
-                           self.figure.dpi, metadata=metadata)
+            with cbook.open_file_cm(filename_or_obj, "wb") as fh:
+                _png.write_png(renderer._renderer, fh,
+                               self.figure.dpi, metadata=metadata)
         finally:
-            if close:
-                filename_or_obj.close()
             renderer.dpi = original_dpi
 
     def print_to_buffer(self):
diff --git a/lib/matplotlib/backends/backend_pdf.py b/lib/matplotlib/backends/backend_pdf.py
@@ -27,7 +27,7 @@
 
 import numpy as np
 
-from matplotlib import __version__, rcParams
+from matplotlib import cbook, __version__, rcParams
 from matplotlib._pylab_helpers import Gcf
 from matplotlib.backend_bases import (
     _Backend, FigureCanvasBase, FigureManagerBase, GraphicsContextBase,
@@ -435,9 +435,8 @@ def __init__(self, filename, metadata=None):
         self.passed_in_file_object = False
         self.original_file_like = None
         self.tell_base = 0
-        if isinstance(filename, six.string_types):
-            fh = open(filename, 'wb')
-        elif is_writable_file_like(filename):
+        fh, opened = cbook.to_filehandle(filename, "wb", return_opened=True)
+        if not opened:
             try:
                 self.tell_base = filename.tell()
             except IOError:
@@ -446,8 +445,6 @@ def __init__(self, filename, metadata=None):
             else:
                 fh = filename
                 self.passed_in_file_object = True
-        else:
-            raise ValueError("filename must be a path or a file-like object")
 
         self._core14fontdir = os.path.join(
             rcParams['datapath'], 'fonts', 'pdfcorefonts')
diff --git a/lib/matplotlib/backends/backend_ps.py b/lib/matplotlib/backends/backend_ps.py
@@ -13,7 +13,7 @@
 import logging
 
 from tempfile import mkstemp
-from matplotlib import __version__, rcParams, checkdep_ghostscript
+from matplotlib import cbook, __version__, rcParams, checkdep_ghostscript
 from matplotlib.afm import AFM
 from matplotlib.backend_bases import (
     _Backend, FigureCanvasBase, FigureManagerBase, GraphicsContextBase,
@@ -982,9 +982,10 @@ def _print_figure(self, outfile, format, dpi=72, facecolor='w', edgecolor='w',
         the key 'Creator' is used.
         """
         isEPSF = format == 'eps'
-        passed_in_file_object = False
-        if isinstance(outfile, six.string_types):
-            title = outfile
+        if isinstance(outfile,
+                      (six.string_types, getattr(os, "PathLike", ()),)):
+            outfile = title = getattr(os, "fspath", lambda obj: obj)(outfile)
+            passed_in_file_object = False
         elif is_writable_file_like(outfile):
             title = None
             passed_in_file_object = True
diff --git a/lib/matplotlib/backends/backend_svg.py b/lib/matplotlib/backends/backend_svg.py
@@ -18,11 +18,10 @@
 
 import numpy as np
 
-from matplotlib import __version__, rcParams
+from matplotlib import cbook, __version__, rcParams
 from matplotlib.backend_bases import (
      _Backend, FigureCanvasBase, FigureManagerBase, RendererBase)
 from matplotlib.backends.backend_mixed import MixedModeRenderer
-from matplotlib.cbook import is_writable_file_like
 from matplotlib.colors import rgb2hex
 from matplotlib.font_manager import findfont, get_font
 from matplotlib.ft2font import LOAD_NO_HINTING
@@ -1189,61 +1188,49 @@ class FigureCanvasSVG(FigureCanvasBase):
     fixed_dpi = 72
 
     def print_svg(self, filename, *args, **kwargs):
-        if isinstance(filename, six.string_types):
-            with io.open(filename, 'w', encoding='utf-8') as svgwriter:
-                return self._print_svg(filename, svgwriter, **kwargs)
+        with cbook.open_file_cm(filename, "w", encoding="utf-8") as fh:
 
-        if not is_writable_file_like(filename):
-            raise ValueError("filename must be a path or a file-like object")
+            filename = getattr(fh, 'name', '')
+            if not isinstance(filename, six.string_types):
+                filename = ''
 
-        svgwriter = filename
-        filename = getattr(svgwriter, 'name', '')
-        if not isinstance(filename, six.string_types):
-            filename = ''
-
-        if not isinstance(svgwriter, io.TextIOBase):
-            if six.PY3:
-                svgwriter = io.TextIOWrapper(svgwriter, 'utf-8')
+            if cbook.file_requires_unicode(fh):
+                detach = False
             else:
-                svgwriter = codecs.getwriter('utf-8')(svgwriter)
-            detach = True
-        else:
-            detach = False
+                if six.PY3:
+                    fh = io.TextIOWrapper(fh, 'utf-8')
+                else:
+                    fh = codecs.getwriter('utf-8')(fh)
+                detach = True
 
-        result = self._print_svg(filename, svgwriter, **kwargs)
+            result = self._print_svg(filename, fh, **kwargs)
 
-        # Detach underlying stream from wrapper so that it remains open in the
-        # caller.
-        if detach:
-            if six.PY3:
-                svgwriter.detach()
-            else:
-                svgwriter.reset()
-                svgwriter.stream = io.BytesIO()
+            # Detach underlying stream from wrapper so that it remains open in
+            # the caller.
+            if detach:
+                if six.PY3:
+                    fh.detach()
+                else:
+                    fh.reset()
+                    fh.stream = io.BytesIO()
 
         return result
 
     def print_svgz(self, filename, *args, **kwargs):
-        if isinstance(filename, six.string_types):
-            options = dict(filename=filename)
-        elif is_writable_file_like(filename):
-            options = dict(fileobj=filename)
-        else:
-            raise ValueError("filename must be a path or a file-like object")
-
-        with gzip.GzipFile(mode='w', **options) as gzipwriter:
+        with cbook.open_file_cm(filename, "wb") as fh, \
+                gzip.GzipFile(mode='w', fileobj=fh) as gzipwriter:
             return self.print_svg(gzipwriter)
 
-    def _print_svg(self, filename, svgwriter, **kwargs):
+    def _print_svg(self, filename, fh, **kwargs):
         image_dpi = kwargs.pop("dpi", 72)
         self.figure.set_dpi(72.0)
         width, height = self.figure.get_size_inches()
-        w, h = width*72, height*72
+        w, h = width * 72, height * 72
 
         _bbox_inches_restore = kwargs.pop("bbox_inches_restore", None)
         renderer = MixedModeRenderer(
-            self.figure,
-            width, height, image_dpi, RendererSVG(w, h, svgwriter, filename, image_dpi),
+            self.figure, width, height, image_dpi,
+            RendererSVG(w, h, fh, filename, image_dpi),
             bbox_inches_restore=_bbox_inches_restore)
 
         self.figure.draw(renderer)
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -10,7 +10,9 @@
 
 import six
 from six.moves import xrange, zip
+import bz2
 import collections
+import contextlib
 import datetime
 import errno
 import functools
@@ -584,12 +586,16 @@ def is_numlike(obj):
     return isinstance(obj, (numbers.Number, np.number))
 
 
-def to_filehandle(fname, flag='rU', return_opened=False):
+def to_filehandle(fname, flag='rU', return_opened=False, encoding=None):
     """
-    *fname* can be a filename or a file handle.  Support for gzipped
+    *fname* can be an `os.PathLike` or a file handle.  Support for gzipped
     files is automatic, if the filename ends in .gz.  *flag* is a
     read/write flag for :func:`file`
     """
+    if hasattr(os, "PathLike") and isinstance(fname, os.PathLike):
+        return to_filehandle(
+            os.fspath(fname),
+            flag=flag, return_opened=return_opened, encoding=encoding)
     if isinstance(fname, six.string_types):
         if fname.endswith('.gz'):
             # get rid of 'U' in flag for gzipped files.
@@ -598,21 +604,31 @@ def to_filehandle(fname, flag='rU', return_opened=False):
         elif fname.endswith('.bz2'):
             # get rid of 'U' in flag for bz2 files
             flag = flag.replace('U', '')
-            import bz2
             fh = bz2.BZ2File(fname, flag)
         else:
-            fh = open(fname, flag)
+            fh = io.open(fname, flag, encoding=encoding)
         opened = True
     elif hasattr(fname, 'seek'):
         fh = fname
         opened = False
     else:
-        raise ValueError('fname must be a string or file handle')
+        raise ValueError('fname must be a PathLike or file handle')
     if return_opened:
         return fh, opened
     return fh
 
 
+@contextlib.contextmanager
+def open_file_cm(path_or_file, mode="r", encoding=None):
+    r"""Pass through file objects and context-manage `~.PathLike`\s."""
+    fh, opened = to_filehandle(path_or_file, mode, True, encoding)
+    if opened:
+        with fh:
+            yield fh
+    else:
+        yield fh
+
+
 def is_scalar_or_string(val):
     """Return whether the given object is a scalar or string like."""
     return isinstance(val, six.string_types) or not iterable(val)
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -1360,11 +1360,12 @@ def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
     """
     from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
     from matplotlib.figure import Figure
-
-    # Fast path for saving to PNG
-    if (format == 'png' or format is None or
-            isinstance(fname, six.string_types) and
-            fname.lower().endswith('.png')):
+    if isinstance(fname, getattr(os, "PathLike", ())):
+        fname = os.fspath(fname)
+    if (format == 'png'
+        or (format is None
+            and isinstance(fname, six.string_types)
+            and fname.lower().endswith('.png'))):
         image = AxesImage(None, cmap=cmap, origin=origin)
         image.set_data(arr)
         image.set_clim(vmin, vmax)
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -111,22 +111,26 @@ def isAvailable(self):
 
 
 WRITER_OUTPUT = [
-    ('ffmpeg', 'mp4'),
-    ('ffmpeg_file', 'mp4'),
-    ('avconv', 'mp4'),
-    ('avconv_file', 'mp4'),
-    ('imagemagick', 'gif'),
-    ('imagemagick_file', 'gif'),
-    ('html', 'html'),
-    ('null', 'null')
+    ('ffmpeg', 'movie.mp4'),
+    ('ffmpeg_file', 'movie.mp4'),
+    ('avconv', 'movie.mp4'),
+    ('avconv_file', 'movie.mp4'),
+    ('imagemagick', 'movie.gif'),
+    ('imagemagick_file', 'movie.gif'),
+    ('html', 'movie.html'),
+    ('null', 'movie.null')
 ]
+if sys.version_info >= (3, 6):
+    from pathlib import Path
+    WRITER_OUTPUT += [
+        (writer, Path(output)) for writer, output in WRITER_OUTPUT]
 
 
 # Smoke test for saving animations.  In the future, we should probably
 # design more sophisticated tests which compare resulting frames a-la
 # matplotlib.testing.image_comparison
-@pytest.mark.parametrize('writer, extension', WRITER_OUTPUT)
-def test_save_animation_smoketest(tmpdir, writer, extension):
+@pytest.mark.parametrize('writer, output', WRITER_OUTPUT)
+def test_save_animation_smoketest(tmpdir, writer, output):
     try:
         # for ImageMagick the rcparams must be patched to account for
         # 'convert' being a built in MS tool, not the imagemagick
@@ -165,8 +169,8 @@ def animate(i):
     with tmpdir.as_cwd():
         anim = animation.FuncAnimation(fig, animate, init_func=init, frames=5)
         try:
-            anim.save('movie.' + extension, fps=30, writer=writer, bitrate=500,
-                      dpi=dpi, codec=codec)
+            anim.save(output, fps=30, writer=writer, bitrate=500, dpi=dpi,
+                      codec=codec)
         except UnicodeDecodeError:
             pytest.xfail("There can be errors in the numpy import stack, "
                          "see issues #1891 and #2679")
diff --git a/lib/matplotlib/tests/test_backend_pdf.py b/lib/matplotlib/tests/test_backend_pdf.py
@@ -135,6 +135,13 @@ def test_composite_image():
         assert len(pdf._file._images) == 2
 
 
+@pytest.mark.skipif(sys.version_info < (3, 6), reason="requires Python 3.6+")
+def test_pdfpages_fspath():
+    from pathlib import Path
+    with PdfPages(Path(os.devnull)) as pdf:
+        pdf.savefig(plt.figure())
+
+
 def test_source_date_epoch():
     """Test SOURCE_DATE_EPOCH support for PDF output"""
     _determinism_source_date_epoch("pdf", b"/CreationDate (D:20000101000000Z)")
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -1,13 +1,16 @@
 from __future__ import absolute_import, division, print_function
 
+import os
+import sys
+import warnings
+
 from matplotlib import rcParams
 from matplotlib.testing.decorators import image_comparison
 from matplotlib.axes import Axes
 from matplotlib.ticker import AutoMinorLocator, FixedFormatter
 import matplotlib.pyplot as plt
 import matplotlib.dates as mdates
 import numpy as np
-import warnings
 import pytest
 
 
@@ -321,7 +324,7 @@ def test_subplots_shareax_loglabels():
 
 
 def test_savefig():
-    fig, ax = plt.subplots()
+    fig = plt.figure()
     msg = "savefig() takes 2 positional arguments but 3 were given"
     with pytest.raises(TypeError, message=msg):
         fig.savefig("fname1.png", "fname2.png")
@@ -330,3 +333,10 @@ def test_savefig():
 def test_figure_repr():
     fig = plt.figure(figsize=(10, 20), dpi=10)
     assert repr(fig) == "<Figure size 100x200 with 0 Axes>"
+
+
+@pytest.mark.skipif(sys.version_info < (3, 6), reason="requires Python 3.6+")
+@pytest.mark.parametrize("fmt", ["png", "pdf", "ps", "eps", "svg"])
+def test_fspath(fmt):
+    from pathlib import Path
+    plt.savefig(Path(os.devnull), format=fmt)
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
