67 строки
1.7 KiB
Python
67 строки
1.7 KiB
Python
import os
|
|
import shutil
|
|
import tempfile
|
|
import atexit
|
|
import stat
|
|
|
|
def try_delete(filename):
|
|
try:
|
|
os.unlink(filename)
|
|
except:
|
|
pass
|
|
if not os.path.exists(filename): return
|
|
try:
|
|
shutil.rmtree(filename, ignore_errors=True)
|
|
except:
|
|
pass
|
|
if not os.path.exists(filename): return
|
|
try:
|
|
os.chmod(filename, os.stat(filename).st_mode | stat.S_IWRITE)
|
|
def remove_readonly_and_try_again(func, path, exc_info):
|
|
if not (os.stat(path).st_mode & stat.S_IWRITE):
|
|
os.chmod(path, os.stat(path).st_mode | stat.S_IWRITE)
|
|
func(path)
|
|
else:
|
|
raise
|
|
shutil.rmtree(filename, onerror=remove_readonly_and_try_again)
|
|
except:
|
|
pass
|
|
|
|
class TempFiles:
|
|
def __init__(self, tmp, save_debug_files=False):
|
|
self.tmp = tmp
|
|
self.save_debug_files = save_debug_files
|
|
self.to_clean = []
|
|
|
|
atexit.register(self.clean)
|
|
|
|
def note(self, filename):
|
|
self.to_clean.append(filename)
|
|
|
|
def get(self, suffix):
|
|
"""Returns a named temp file with the given prefix."""
|
|
named_file = tempfile.NamedTemporaryFile(dir=self.tmp, suffix=suffix, delete=False)
|
|
self.note(named_file.name)
|
|
return named_file
|
|
|
|
def get_dir(self):
|
|
"""Returns a named temp file with the given prefix."""
|
|
directory = tempfile.mkdtemp(dir=self.tmp)
|
|
self.note(directory)
|
|
return directory
|
|
|
|
def clean(self):
|
|
if self.save_debug_files:
|
|
import sys
|
|
print >> sys.stderr, 'not cleaning up temp files since in debug-save mode, see them in %s' % (self.tmp,)
|
|
return
|
|
for filename in self.to_clean:
|
|
try_delete(filename)
|
|
self.to_clean = []
|
|
|
|
def run_and_clean(self, func):
|
|
try:
|
|
return func()
|
|
finally:
|
|
self.clean()
|