8000 Use tempfile.mkdtemp for animation.FileMovieWriter. · matplotlib/matplotlib@f10180a · GitHub
[go: up one dir, main page]
Skip to content

Commit f10180a

Browse files
anntzeranntzer
committed
Use tempfile.mkdtemp for animation.FileMovieWriter.
Currently, `animation.FileMovieWriter` (e.g. `FFMpegFileWriter`) creates temporaries in the current folder, without even checking whether the file existed before. This is obviously risky (e.g., create a file named `_tmp0000000.png` in the current folder before running `moviewriter.py` from the docs examples and it will be overwritten. Moveover, temporaries are not cleaned up if an exception is raised (e.g., KeyboardInterrupt). This is relatively easy to fix using `tempfile.mkdtemp`, although I had to modify `finish` so that it doesn't call `cleanup` by itself (obviously, this can be fixed by keeping a `cleaned_up` attribute but it feels like a minor API change is preferable; moreover the context-manager API stays the same). Only basic, manual testing done.
1 parent 2c11590 commit f10180a

File tree

1 file changed

+35
-18
lines changed

1 file changed

+35
-18
lines changed

‎lib/matplotlib/animation.py

Lines changed: 35 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,12 @@
2323
import six
2424
from six.moves import xrange, zip
2525

26-
import sys
27-
import itertools
2826
import contextlib
27+
import itertools
28+
import os
29+
import shutil
30+
import sys
31+
import tempfile
2932
from matplotlib.cbook import iterable, is_string_like
3033
from matplotlib.compat import subprocess
3134
from matplotlib import verbose
@@ -173,8 +176,11 @@ def saving(self, *args):
173176
'''
174177
# This particular sequence is what contextlib.contextmanager wants
175178
self.setup(*args)
176-
yield
177-
self.finish()
179+
try:
180+
yield
181+
self.finish()
182+
finally:
183+
self.cleanup()
178184

179185
def _run(self):
180186
# Uses subprocess to call the program for assembling frames into a
@@ -193,7 +199,6 @@ def _run(self):
193199

194200
def finish(self):
195201
'Finish any processing for writing the movie.'
196-
self.cleanup()
197202

198203
def grab_frame(self, **savefig_kwargs):
199204
'''
@@ -225,11 +230,12 @@ def _args(self):
225230

226231
def cleanup(self):
227232
'Clean-up and collect the process used to write the movie file.'
228-
out, err = self._proc.communicate()
229-
verbose.report('MovieWriter -- '
230-
'Command stdout:\n%s' % out, level='debug')
231-
verbose.report('MovieWriter -- '
232-
'Command stderr:\n%s' % err, level='debug')
233+
if hasattr(self, "_proc"):
234+
out, err = self._proc.communicate()
235+
verbose.report('MovieWriter -- Command stdout:\n%s' % out,
236+
level='debug')
237+
verbose.report('MovieWriter -- Command stderr:\n%s' % err,
238+
level='debug')
233239

234240
@classmethod
235241
def bin_path(cls):
@@ -263,7 +269,8 @@ def __init__(self, *args, **kwargs):
263269
MovieWriter.__init__(self, *args, **kwargs)
264270
self.frame_format = rcParams['animation.frame_format']
265271

266-
def setup(self, fig, outfile, dpi, frame_prefix='_tmp', clear_temp=True):
272+
def setup(self, fig, outfile, dpi, frame_prefix='_tmp', clear_temp=True,
273+
tmpdir=None):
267274
'''
268275
Perform setup for writing the movie file.
269276
@@ -280,15 +287,22 @@ def setup(self, fig, outfile, dpi, frame_prefix='_tmp', clear_temp=True):
280287
clear_temp: bool
281288
Specifies whether the temporary files should be deleted after
282289
the movie is written. (Useful for debugging.) Defaults to True.
290+
tmpdir: string, optional
291+
Path to temporary directory.
283292
'''
284293
self.fig = fig
285294
self.outfile = outfile
286295
self.dpi = dpi
287296
self.clear_temp = clear_temp
288297
self.temp_prefix = frame_prefix
289298
self._frame_counter = 0 # used for generating sequential file names
290-
self._temp_names = list()
291-
self.fname_format_str = '%s%%07d.%s'
299+
if tmpdir is None:
300+
self._tmpdir = self._temp_names = tempfile.mkdtemp()
301+
else:
302+
self._tmpdir = tmpdir
303+
self._temp_names = list()
304+
self.fname_format_str = os.path.join(
305+
self._tmpdir.replace('%', '%%'), '%s%%07d.%s')
292306

293307
@property
294308
def frame_format(self):
@@ -316,7 +330,8 @@ def _frame_sink(self):
316330
fname = self._base_temp_name() % self._frame_counter
317331

318332
# Save the filename so we can delete it later if necessary
319-
self._temp_names.append(fname)
333+
if isinstance(self._temp_names, list):
334+
self._temp_names.append(fname)
320335
verbose.report(
321336
'FileMovieWriter.frame_sink: saving frame %d to fname=%s' %
322337
(self._frame_counter, fname),
@@ -355,7 +370,7 @@ def finish(self):
355370
# Call run here now that all frame grabbing is done. All temp files
356371
# are available to be assembled.
357372
self._run()
358-
MovieWriter.finish(self) # Will call clean-up
373+
MovieWriter.finish(self)
359374

360375
# Check error code for creating file here, since we just run
361376
# the process here, rather than having an open pipe.
@@ -369,13 +384,15 @@ def cleanup(self):
369384

370385
#Delete temporary files
371386
if self.clear_temp:
372-
import os
373387
verbose.report(
374388
'MovieWriter: clearing temporary fnames=%s' %
375389
str(self._temp_names),
376390
level='debug')
377-
for fname in self._temp_names:
378-
os.remove(fname)
391+
if isinstance(self._temp_names, list):
392+
for fname in self._temp_names:
393+
os.remove(fname)
394+
else:
395+
shutil.rmtree(self._temp_names)
379396

380397

381398
# Base class of ffmpeg information. Has the config keys and the common set

0 commit comments

Comments
 (0)
0