blob: 31c79293229c762110dbf38d1aaf95be4ffd6b66 [file] [log] [blame]
"""
Autotest tempfile wrapper for mkstemp (known as tempfile here) and
mkdtemp (known as tempdir).
This wrapper provides a mechanism to clean up temporary files/dirs once they
are no longer need.
Files/Dirs will have a unique_id prepended to the suffix and a
_autotmp_ tag appended to the prefix.
It is required that the unique_id param is supplied when a temp dir/file is
created.
"""
import shutil, os, logging
import tempfile as module_tempfile
_TEMPLATE = '_autotmp_'
class tempfile(object):
"""
A wrapper for tempfile.mkstemp
@param unique_id: required, a unique string to help identify what
part of code created the tempfile.
@param auto_clean: automatically delete the temporary file in destructor.
@var name: The name of the temporary file.
@var fd: the file descriptor of the temporary file that was created.
@return a tempfile object
example usage:
t = autotemp.tempfile(unique_id='fig')
t.name # name of file
t.fd # file descriptor
t.fo # file object
t.clean() # clean up after yourself
"""
def __init__(self, unique_id, suffix='', prefix='', dir=None,
text=False, auto_clean=True):
self.auto_clean = auto_clean
suffix = unique_id + suffix
prefix = prefix + _TEMPLATE
self.fd, self.name = module_tempfile.mkstemp(suffix=suffix,
prefix=prefix,
dir=dir, text=text)
self.fo = os.fdopen(self.fd)
def clean(self):
"""
Remove the temporary file that was created.
This is also called by the destructor.
"""
if self.fo:
self.fo.close()
if self.name and os.path.exists(self.name):
os.remove(self.name)
self.fd = self.fo = self.name = None
def __del__(self):
try:
if self.name is not None and self.auto_clean:
logging.debug('Auto-cleaning %s', self.name)
self.clean()
except:
try:
msg = 'An exception occurred while calling the destructor'
logging.exception(msg)
except:
pass
class tempdir(object):
"""
A wrapper for tempfile.mkdtemp
@var name: The name of the temporary dir.
@return A tempdir object
example usage:
b = autotemp.tempdir(unique_id='exemdir')
b.name # your directory
b.clean() # clean up after yourself
"""
def __init__(self, suffix='', unique_id='', prefix='', dir=None,
auto_clean=True):
"""
Initialize temp directory.
@param suffix: suffix for dir.
@param prefix: prefix for dir. Defaults to '_autotmp'.
@param unique_id: unique id of tempdir.
@param dir: parent directory of the tempdir. Defaults to /tmp.
@param auto_clean: automatically clean up the tempdir in destructor.
eg: autotemp.tempdir(suffix='suffix', unique_id='123', prefix='prefix')
creates a dir like '/tmp/prefix_autotmp_<random hash>123suffix'
"""
self.auto_clean = auto_clean
suffix = unique_id + suffix
prefix = prefix + _TEMPLATE
self.name = module_tempfile.mkdtemp(suffix=suffix,
prefix=prefix, dir=dir)
def clean(self):
"""
Remove the temporary dir that was created.
This is also called by the destructor.
"""
if self.name and os.path.exists(self.name):
shutil.rmtree(self.name)
self.name = None
def __del__(self):
try:
if self.name and self.auto_clean:
logging.debug('Auto-cleaning %s', self.name)
self.clean()
except:
try:
msg = 'An exception occurred while calling the destructor'
logging.exception(msg)
except:
pass
class dummy_dir(object):
"""A dummy object representing a directory with a name.
Only used for compat with the tmpdir, in cases where we wish to
reuse a dir with the same interface but not to delete it after
we're done using it.
"""
def __init__(self, name):
"""Initialize the dummy_dir object.
@param name: Path the the directory.
"""
self.name = name