Skip to content

Commit 89a130f

Browse files
anntzeranntzer
committed
Support PathLike inputs.
1 parent f6ebbc3 commit 89a130f

12 files changed

+126
-94
lines changed

‎doc/users/next_whats_new/2018-01-10-AL-pathlike.rst

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
PathLike support
2+
````````````````
3+
4+
The following functions now support `~.PathLike` inputs on Python 3.6+:
5+
On Python 3.6+, `~matplotlib.pyplot.savefig`, `~matplotlib.pyplot.imsave`,
6+
`~matplotlib.pyplot.imread`, and animation writers now accept `~os.PathLike`\s
7+
as input.

‎lib/matplotlib/animation.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -358,8 +358,7 @@ def _run(self):
358358
# from a few configuration options.
359359
command = self._args()
360360
output = subprocess.PIPE
361-
_log.info('MovieWriter.run: running command: %s',
362-
' '.join(command))
361+
_log.info('MovieWriter.run: running command: %s', command)
363362
self._proc = subprocess.Popen(command, shell=False,
364363
stdout=output, stderr=output,
365364
stdin=subprocess.PIPE,
@@ -847,12 +846,13 @@ def __init__(self, fps=30, codec=None, bitrate=None, extra_args=None,
847846
extra_args, metadata)
848847

849848
def setup(self, fig, outfile, dpi, frame_dir=None):
850-
if os.path.splitext(outfile)[-1] not in ['.html', '.htm']:
849+
root, ext = os.path.splitext(outfile)
850+
if ext not in ['.html', '.htm']:
851851
raise ValueError("outfile must be *.htm or *.html")
852852

853853
if not self.embed_frames:
854854
if frame_dir is None:
855-
frame_dir = outfile.rstrip('.html') + '_frames'
855+
frame_dir = root + '_frames'
856856
if not os.path.exists(frame_dir):
857857
os.makedirs(frame_dir)
858858
frame_prefix = os.path.join(frame_dir, 'frame')

‎lib/matplotlib/backends/backend_agg.py

+3-9
Original file line numberDiff line numberDiff line change
@@ -510,11 +510,6 @@ def print_png(self, filename_or_obj, *args, **kwargs):
510510
renderer = self.get_renderer()
511511
original_dpi = renderer.dpi
512512
renderer.dpi = self.figure.dpi
513-
if isinstance(filename_or_obj, six.string_types):
514-
filename_or_obj = open(filename_or_obj, 'wb')
515-
close = True
516-
else:
517-
close = False
518513

519514
version_str = 'matplotlib version ' + __version__ + \
520515
', http://matplotlib.org/'
@@ -524,11 +519,10 @@ def print_png(self, filename_or_obj, *args, **kwargs):
524519
metadata.update(user_metadata)
525520

526521
try:
527-
_png.write_png(renderer._renderer, filename_or_obj,
528-
self.figure.dpi, metadata=metadata)
522+
with cbook.open_file_cm(filename_or_obj, "wb") as file:
523+
_png.write_png(renderer._renderer, file,
524+
self.figure.dpi, metadata=metadata)
529525
finally:
530-
if close:
531-
filename_or_obj.close()
532526
renderer.dpi = original_dpi
533527

534528
def print_to_buffer(self):

‎lib/matplotlib/backends/backend_pdf.py

+3-6
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727

2828
import numpy as np
2929

30-
from matplotlib import __version__, rcParams
30+
from matplotlib import cbook, __version__, rcParams
3131
from matplotlib._pylab_helpers import Gcf
3232
from matplotlib.backend_bases import (
3333
_Backend, FigureCanvasBase, FigureManagerBase, GraphicsContextBase,
@@ -435,9 +435,8 @@ def __init__(self, filename, metadata=None):
435435
self.passed_in_file_object = False
436436
self.original_file_like = None
437437
self.tell_base = 0
438-
if isinstance(filename, six.string_types):
439-
fh = open(filename, 'wb')
440-
elif is_writable_file_like(filename):
438+
fh, opened = cbook.to_filehandle(filename, "wb", return_opened=True)
439+
if not opened:
441440
try:
442441
self.tell_base = filename.tell()
443442
except IOError:
@@ -446,8 +445,6 @@ def __init__(self, filename, metadata=None):
446445
else:
447446
fh = filename
448447
self.passed_in_file_object = True
449-
else:
450-
raise ValueError("filename must be a path or a file-like object")
451448

452449
self._core14fontdir = os.path.join(
453450
rcParams['datapath'], 'fonts', 'pdfcorefonts')

‎lib/matplotlib/backends/backend_ps.py

+5-4
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import logging
1414

1515
from tempfile import mkstemp
16-
from matplotlib import __version__, rcParams, checkdep_ghostscript
16+
from matplotlib import cbook, __version__, rcParams, checkdep_ghostscript
1717
from matplotlib.afm import AFM
1818
from matplotlib.backend_bases import (
1919
_Backend, FigureCanvasBase, FigureManagerBase, GraphicsContextBase,
@@ -961,9 +961,10 @@ def _print_figure(self, outfile, format, dpi=72, facecolor='w', edgecolor='w',
961961
the key 'Creator' is used.
962962
"""
963963
isEPSF = format == 'eps'
964-
passed_in_file_object = False
965-
if isinstance(outfile, six.string_types):
966-
title = outfile
964+
if isinstance(outfile,
965+
(six.string_types, getattr(os, "PathLike", ()),)):
966+
outfile = title = getattr(os, "fspath", lambda obj: obj)(outfile)
967+
passed_in_file_object = False
967968
elif is_writable_file_like(outfile):
968969
title = None
969970
passed_in_file_object = True

‎lib/matplotlib/backends/backend_svg.py

+25-38
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,10 @@
1818

1919
import numpy as np
2020

21-
from matplotlib import __version__, rcParams
21+
from matplotlib import cbook, __version__, rcParams
2222
from matplotlib.backend_bases import (
2323
_Backend, FigureCanvasBase, FigureManagerBase, RendererBase)
2424
from matplotlib.backends.backend_mixed import MixedModeRenderer
25-
from matplotlib.cbook import is_writable_file_like
2625
from matplotlib.colors import rgb2hex
2726
from matplotlib.font_manager import findfont, get_font
2827
from matplotlib.ft2font import LOAD_NO_HINTING
@@ -1189,52 +1188,40 @@ class FigureCanvasSVG(FigureCanvasBase):
11891188
fixed_dpi = 72
11901189

11911190
def print_svg(self, filename, *args, **kwargs):
1192-
if isinstance(filename, six.string_types):
1193-
with io.open(filename, 'w', encoding='utf-8') as svgwriter:
1194-
return self._print_svg(filename, svgwriter, **kwargs)
1191+
with cbook.open_file_cm(filename, "w", encoding="utf-8") as file:
11951192

1196-
if not is_writable_file_like(filename):
1197-
raise ValueError("filename must be a path or a file-like object")
1193+
filename = getattr(file, 'name', '')
1194+
if not isinstance(filename, six.string_types):
1195+
filename = ''
11981196

1199-
svgwriter = filename
1200-
filename = getattr(svgwriter, 'name', '')
1201-
if not isinstance(filename, six.string_types):
1202-
filename = ''
1203-
1204-
if not isinstance(svgwriter, io.TextIOBase):
1205-
if six.PY3:
1206-
svgwriter = io.TextIOWrapper(svgwriter, 'utf-8')
1197+
if not isinstance(file, io.TextIOBase):
1198+
if six.PY3:
1199+
file = io.TextIOWrapper(file, 'utf-8')
1200+
else:
1201+
file = codecs.getwriter('utf-8')(file)
1202+
detach = True
12071203
else:
1208-
svgwriter = codecs.getwriter('utf-8')(svgwriter)
1209-
detach = True
1210-
else:
1211-
detach = False
1204+
detach = False
12121205

1213-
result = self._print_svg(filename, svgwriter, **kwargs)
1206+
result = self._print_svg(filename, file, **kwargs)
12141207

1215-
# Detach underlying stream from wrapper so that it remains open in the
1216-
# caller.
1217-
if detach:
1218-
if six.PY3:
1219-
svgwriter.detach()
1220-
else:
1221-
svgwriter.reset()
1222-
svgwriter.stream = io.BytesIO()
1208+
# Detach underlying stream from wrapper so that it remains open in
1209+
# the caller.
1210+
if detach:
1211+
if six.PY3:
1212+
file.detach()
1213+
else:
1214+
file.reset()
1215+
file.stream = io.BytesIO()
12231216

12241217
return result
12251218

12261219
def print_svgz(self, filename, *args, **kwargs):
1227-
if isinstance(filename, six.string_types):
1228-
options = dict(filename=filename)
1229-
elif is_writable_file_like(filename):
1230-
options = dict(fileobj=filename)
1231-
else:
1232-
raise ValueError("filename must be a path or a file-like object")
1233-
1234-
with gzip.GzipFile(mode='w', **options) as gzipwriter:
1220+
with cbook.open_file_cm(filename, "wb") as file, \
1221+
gzip.GzipFile(mode='w', fileobj=file) as gzipwriter:
12351222
return self.print_svg(gzipwriter)
12361223

1237-
def _print_svg(self, filename, svgwriter, **kwargs):
1224+
def _print_svg(self, filename, file, **kwargs):
12381225
image_dpi = kwargs.pop("dpi", 72)
12391226
self.figure.set_dpi(72.0)
12401227
width, height = self.figure.get_size_inches()
@@ -1243,7 +1230,7 @@ def _print_svg(self, filename, svgwriter, **kwargs):
12431230
_bbox_inches_restore = kwargs.pop("bbox_inches_restore", None)
12441231
renderer = MixedModeRenderer(
12451232
self.figure,
1246-
width, height, image_dpi, RendererSVG(w, h, svgwriter, filename, image_dpi),
1233+
width, height, image_dpi, RendererSVG(w, h, file, filename, image_dpi),
12471234
bbox_inches_restore=_bbox_inches_restore)
12481235

12491236
self.figure.draw(renderer)

‎lib/matplotlib/cbook/__init__.py

+20-4
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import six
1212
from six.moves import xrange, zip
1313
import collections
14+
import contextlib
1415
import datetime
1516
import errno
1617
import functools
@@ -584,12 +585,16 @@ def is_numlike(obj):
584585
return isinstance(obj, (numbers.Number, np.number))
585586

586587

587-
def to_filehandle(fname, flag='rU', return_opened=False):
588+
def to_filehandle(fname, flag='rU', return_opened=False, encoding=None):
588589
"""
589-
*fname* can be a filename or a file handle. Support for gzipped
590+
*fname* can be an `os.PathLike` or a file handle. Support for gzipped
590591
files is automatic, if the filename ends in .gz. *flag* is a
591592
read/write flag for :func:`file`
592593
"""
594+
if hasattr(os, "PathLike") and isinstance(fname, os.PathLike):
595+
return to_filehandle(
596+
os.fspath(fname),
597+
flag=flag, return_opened=return_opened, encoding=encoding)
593598
if isinstance(fname, six.string_types):
594599
if fname.endswith('.gz'):
595600
# get rid of 'U' in flag for gzipped files.
@@ -601,18 +606,29 @@ def to_filehandle(fname, flag='rU', return_opened=False):
601606
import bz2
602607
fh = bz2.BZ2File(fname, flag)
603608
else:
604-
fh = open(fname, flag)
609+
fh = open(fname, flag, encoding=encoding)
605610
opened = True
606611
elif hasattr(fname, 'seek'):
607612
fh = fname
608613
opened = False
609614
else:
610-
raise ValueError('fname must be a string or file handle')
615+
raise ValueError('fname must be a PathLike or file handle')
611616
if return_opened:
612617
return fh, opened
613618
return fh
614619

615620

621+
@contextlib.contextmanager
622+
def open_file_cm(path_or_file, mode="r", encoding=None):
623+
r"""Pass through file objects and context-manage `~.PathLike`\s."""
624+
fh, opened = to_filehandle(path_or_file, mode, True, encoding=encoding)
625+
if opened:
626+
with fh:
627+
yield fh
628+
else:
629+
yield fh
630+
631+
616632
def is_scalar_or_string(val):
617633
"""Return whether the given object is a scalar or string like."""
618634
return isinstance(val, six.string_types) or not iterable(val)

‎lib/matplotlib/image.py

+5-15
Original file line numberDiff line numberDiff line change
@@ -1349,21 +1349,11 @@ def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
13491349
"""
13501350
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
13511351
from matplotlib.figure import Figure
1352-
1353-
# Fast path for saving to PNG
1354-
if (format == 'png' or format is None or
1355-
isinstance(fname, six.string_types) and
1356-
fname.lower().endswith('.png')):
1357-
image = AxesImage(None, cmap=cmap, origin=origin)
1358-
image.set_data(arr)
1359-
image.set_clim(vmin, vmax)
1360-
image.write_png(fname)
1361-
else:
1362-
fig = Figure(dpi=dpi, frameon=False)
1363-
FigureCanvas(fig)
1364-
fig.figimage(arr, cmap=cmap, vmin=vmin, vmax=vmax, origin=origin,
1365-
resize=True)
1366-
fig.savefig(fname, dpi=dpi, format=format, transparent=True)
1352+
fig = Figure(dpi=dpi, frameon=False)
1353+
FigureCanvas(fig)
1354+
fig.figimage(arr, cmap=cmap, vmin=vmin, vmax=vmax, origin=origin,
1355+
resize=True)
1356+
fig.savefig(fname, dpi=dpi, format=format, transparent=True)
13671357

13681358

13691359
def pil_to_array(pilImage):

‎lib/matplotlib/tests/test_animation.py

+19-12
Original file line numberDiff line numberDiff line change
@@ -111,22 +111,29 @@ def isAvailable(self):
111111

112112

113113
WRITER_OUTPUT = [
114-
('ffmpeg', 'mp4'),
115-
('ffmpeg_file', 'mp4'),
116-
('avconv', 'mp4'),
117-
('avconv_file', 'mp4'),
118-
('imagemagick', 'gif'),
119-
('imagemagick_file', 'gif'),
120-
('html', 'html'),
121-
('null', 'null')
114+
('ffmpeg', 'movie.mp4'),
115+
('ffmpeg_file', 'movie.mp4'),
116+
('avconv', 'movie.mp4'),
117+
('avconv_file', 'movie.mp4'),
118+
('imagemagick', 'movie.gif'),
119+
('imagemagick_file', 'movie.gif'),
120+
('html', 'movie.html'),
121+
('null', 'movie.null')
122122
]
123+
try:
124+
from pathlib import Path
125+
except ImportError:
126+
pass
127+
else:
128+
WRITER_OUTPUT += [
129+
(writer, Path(output)) for writer, output in WRITER_OUTPUT]
123130

124131

125132
# Smoke test for saving animations. In the future, we should probably
126133
# design more sophisticated tests which compare resulting frames a-la
127134
# matplotlib.testing.image_comparison
128-
@pytest.mark.parametrize('writer, extension', WRITER_OUTPUT)
129-
def test_save_animation_smoketest(tmpdir, writer, extension):
135+
@pytest.mark.parametrize('writer, output', WRITER_OUTPUT)
136+
def test_save_animation_smoketest(tmpdir, writer, output):
130137
try:
131138
# for ImageMagick the rcparams must be patched to account for
132139
# 'convert' being a built in MS tool, not the imagemagick
@@ -165,8 +172,8 @@ def animate(i):
165172
with tmpdir.as_cwd():
166173
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=5)
167174
try:
168-
anim.save('movie.' + extension, fps=30, writer=writer, bitrate=500,
169-
dpi=dpi, codec=codec)
175+
anim.save(output, fps=30, writer=writer, bitrate=500, dpi=dpi,
176+
codec=codec)
170177
except UnicodeDecodeError:
171178
pytest.xfail("There can be errors in the numpy import stack, "
172179
"see issues #1891 and #2679")

‎lib/matplotlib/tests/test_backend_pdf.py

+6
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,12 @@ def test_composite_image():
135135
assert len(pdf._file._images) == 2
136136

137137

138+
def test_pdfpages_fspath():
139+
Path = pytest.importorskip("pathlib").Path
140+
with PdfPages(Path(os.devnull)) as pdf:
141+
pdf.savefig(plt.figure())
142+
143+
138144
def test_source_date_epoch():
139145
"""Test SOURCE_DATE_EPOCH support for PDF output"""
140146
_determinism_source_date_epoch("pdf", b"/CreationDate (D:20000101000000Z)")

‎lib/matplotlib/tests/test_figure.py

+11-2
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,16 @@
11
from __future__ import absolute_import, division, print_function
22

3+
import os
4+
import sys
5+
import warnings
6+
37
from matplotlib import rcParams
48
from matplotlib.testing.decorators import image_comparison
59
from matplotlib.axes import Axes
610
from matplotlib.ticker import AutoMinorLocator, FixedFormatter
711
import matplotlib.pyplot as plt
812
import matplotlib.dates as mdates
913
import numpy as np
10-
import warnings
1114
import pytest
1215

1316

@@ -321,7 +324,7 @@ def test_subplots_shareax_loglabels():
321324

322325

323326
def test_savefig():
324-
fig, ax = plt.subplots()
327+
fig = plt.figure()
325328
msg = "savefig() takes 2 positional arguments but 3 were given"
326329
with pytest.raises(TypeError, message=msg):
327330
fig.savefig("fname1.png", "fname2.png")
@@ -330,3 +333,9 @@ def test_savefig():
330333
def test_figure_repr():
331334
fig = plt.figure(figsize=(10, 20), dpi=10)
332335
assert repr(fig) == "<Figure size 100x200 with 0 Axes>"
336+
337+
338+
@pytest.mark.parametrize("fmt", ["png", "pdf", "ps", "eps", "svg"])
339+
def test_fspath(fmt):
340+
Path = pytest.importorskip("pathlib").Path
341+
plt.savefig(Path(os.devnull), format=fmt)

0 commit comments

Comments
 (0)