Permalink
Find file
Fetching contributors…
Cannot retrieve contributors at this time
262 lines (220 sloc) 8.11 KB
# Copyright (c) 2002-2010 Zooko Wilcox-O'Hearn
# This file is part of pyutil; see README.rst for licensing terms.
"""
Futz with files like a pro.
"""
import errno, exceptions, os, stat, tempfile
try:
import bsddb
except ImportError:
DBNoSuchFileError = None
else:
DBNoSuchFileError = bsddb.db.DBNoSuchFileError
# read_file() and write_file() copied from Mark Seaborn's blog post. Please
# read it for rationale:
# http://lackingrhoticity.blogspot.com/2009/12/readfile-and-writefile-in-python.html
def read_file(filename, mode='rb'):
fh = open(filename, mode)
try:
return fh.read()
finally:
fh.close()
def write_file(filename, data, mode='wb'):
fh = open(filename, mode)
try:
fh.write(data)
finally:
fh.close()
# For backwards-compatibility in case someone is using these names. We used to
# have a superkludge in fileutil.py under these names.
def rename(src, dst, tries=4, basedelay=0.1):
return os.rename(src, dst)
def remove(f, tries=4, basedelay=0.1):
return os.remove(f)
def rmdir(f, tries=4, basedelay=0.1):
return os.rmdir(f)
class _Dir(object):
"""
Hold a set of files and subdirs and clean them all up when asked to.
"""
def __init__(self, name, cleanup=True):
self.name = name
self.cleanup = cleanup
self.files = []
self.subdirs = set()
def file(self, fname, mode=None):
"""
Create a file in the tempdir and remember it so as to close() it
before attempting to cleanup the temp dir.
@rtype: file
"""
ffn = os.path.join(self.name, fname)
if mode is not None:
fo = open(ffn, mode)
else:
fo = open(ffn)
self.register_file(fo)
return fo
def subdir(self, dirname):
"""
Create a subdirectory in the tempdir and remember it so as to call
shutdown() on it before attempting to clean up.
@rtype: _Dir instance
"""
ffn = os.path.join(self.name, dirname)
sd = _Dir(ffn, self.cleanup)
self.register_subdir(sd)
make_dirs(sd.name)
return sd
def register_file(self, fileobj):
"""
Remember the file object and call close() on it before attempting to
clean up.
"""
self.files.append(fileobj)
def register_subdir(self, dirobj):
"""
Remember the _Dir object and call shutdown() on it before attempting
to clean up.
"""
self.subdirs.add(dirobj)
def shutdown(self):
if self.cleanup:
for subdir in hasattr(self, 'subdirs') and self.subdirs or []:
subdir.shutdown()
for fileobj in hasattr(self, 'files') and self.files or []:
if DBNoSuchFileError is None:
fileobj.close() # "close()" is idempotent so we don't need to catch exceptions here
else:
try:
fileobj.close()
except DBNoSuchFileError:
# Ah, except that the bsddb module's file-like object (a DB object) has a non-idempotent close...
pass
if hasattr(self, 'name'):
rm_dir(self.name)
def __repr__(self):
return "<%s instance at %x %s>" % (self.__class__.__name__, id(self), self.name)
def __str__(self):
return self.__repr__()
def __del__(self):
try:
self.shutdown()
except:
import traceback
traceback.print_exc()
class NamedTemporaryDirectory(_Dir):
"""
Call tempfile.mkdtemp(), store the name of the dir in self.name, and
rm_dir() when it gets garbage collected or "shutdown()".
Also keep track of file objects for files within the tempdir and call
close() on them before rm_dir(). This is a convenient way to open temp
files within the directory, and it is very helpful on Windows because you
can't delete a directory which contains a file which is currently open.
"""
def __init__(self, cleanup=True, *args, **kwargs):
""" If cleanup, then the directory will be rmrf'ed when the object is shutdown. """
name = tempfile.mkdtemp(*args, **kwargs)
_Dir.__init__(self, name, cleanup)
class ReopenableNamedTemporaryFile:
"""
This uses tempfile.mkstemp() to generate a secure temp file. It then closes
the file, leaving a zero-length file as a placeholder. You can get the
filename with ReopenableNamedTemporaryFile.name. When the
ReopenableNamedTemporaryFile instance is garbage collected or its shutdown()
method is called, it deletes the file.
"""
def __init__(self, *args, **kwargs):
fd, self.name = tempfile.mkstemp(*args, **kwargs)
os.close(fd)
def __repr__(self):
return "<%s instance at %x %s>" % (self.__class__.__name__, id(self), self.name)
def __str__(self):
return self.__repr__()
def __del__(self):
self.shutdown()
def shutdown(self):
remove(self.name)
def make_dirs(dirname, mode=0777):
"""
An idempotent version of os.makedirs(). If the dir already exists, do
nothing and return without raising an exception. If this call creates the
dir, return without raising an exception. If there is an error that
prevents creation or if the directory gets deleted after make_dirs() creates
it and before make_dirs() checks that it exists, raise an exception.
"""
tx = None
try:
os.makedirs(dirname, mode)
except OSError, x:
tx = x
if not os.path.isdir(dirname):
if tx:
raise tx
raise exceptions.IOError, "unknown error prevented creation of directory, or deleted the directory immediately after creation: %s" % dirname # careful not to construct an IOError with a 2-tuple, as that has a special meaning...
def rmtree(dirname):
"""
A threadsafe and idempotent version of shutil.rmtree(). If the dir is
already gone, do nothing and return without raising an exception. If this
call removes the dir, return without raising an exception. If there is an
error that prevents deletion or if the directory gets created again after
rm_dir() deletes it and before rm_dir() checks that it is gone, raise an
exception.
"""
excs = []
try:
os.chmod(dirname, stat.S_IWRITE | stat.S_IEXEC | stat.S_IREAD)
for f in os.listdir(dirname):
fullname = os.path.join(dirname, f)
if os.path.isdir(fullname):
rm_dir(fullname)
else:
remove(fullname)
os.rmdir(dirname)
except EnvironmentError, le:
# Ignore "No such file or directory", collect any other exception.
if (le.args[0] != 2 and le.args[0] != 3) or (le.args[0] != errno.ENOENT):
excs.append(le)
except Exception, le:
excs.append(le)
# Okay, now we've recursively removed everything, ignoring any "No
# such file or directory" errors, and collecting any other errors.
if os.path.exists(dirname):
if len(excs) == 1:
raise excs[0]
if len(excs) == 0:
raise OSError, "Failed to remove dir for unknown reason."
raise OSError, excs
def rm_dir(dirname):
# Renamed to be like shutil.rmtree and unlike rmdir.
return rmtree(dirname)
def remove_if_possible(f):
try:
remove(f)
except EnvironmentError:
pass
def remove_if_present(f):
try:
remove(f)
except EnvironmentError, le:
# Ignore "No such file or directory", re-raise any other exception.
if (le.args[0] != 2 and le.args[0] != 3) or (le.args[0] != errno.ENOENT):
raise
def rmdir_if_possible(f):
try:
rmdir(f)
except EnvironmentError:
pass
def open_or_create(fname, binarymode=True):
try:
f = open(fname, binarymode and "r+b" or "r+")
except EnvironmentError:
f = open(fname, binarymode and "w+b" or "w+")
return f
def du(basedir):
size = 0
for root, dirs, files in os.walk(basedir):
for f in files:
fn = os.path.join(root, f)
size += os.path.getsize(fn)
return size