Home | History | Annotate | Download | only in common_lib 
      1 """
      2 Autotest tempfile wrapper for mkstemp (known as tempfile here) and
      3 mkdtemp (known as tempdir).
      4 
      5 This wrapper provides a mechanism to clean up temporary files/dirs once they
      6 are no longer need.
      7 
      8 Files/Dirs will have a unique_id prepended to the suffix and a
      9 _autotmp_ tag appended to the prefix.
     10 
     11 It is required that the unique_id param is supplied when a temp dir/file is
     12 created.
     13 """
     14 
     15 import shutil, os, logging
     16 import tempfile as module_tempfile
     17 
     18 _TEMPLATE = '_autotmp_'
     19 
     20 
     21 class tempfile(object):
     22     """
     23     A wrapper for tempfile.mkstemp
     24 
     25     @param unique_id: required, a unique string to help identify what
     26                       part of code created the tempfile.
     27     @var name: The name of the temporary file.
     28     @var fd:  the file descriptor of the temporary file that was created.
     29     @return a tempfile object
     30     example usage:
     31         t = autotemp.tempfile(unique_id='fig')
     32         t.name # name of file
     33         t.fd   # file descriptor
     34         t.fo   # file object
     35         t.clean() # clean up after yourself
     36     """
     37     def __init__(self, unique_id, suffix='', prefix='', dir=None,
     38                  text=False):
     39         suffix = unique_id + suffix
     40         prefix = prefix + _TEMPLATE
     41         self.fd, self.name = module_tempfile.mkstemp(suffix=suffix,
     42                                                      prefix=prefix,
     43                                                      dir=dir, text=text)
     44         self.fo = os.fdopen(self.fd)
     45 
     46 
     47     def clean(self):
     48         """
     49         Remove the temporary file that was created.
     50         This is also called by the destructor.
     51         """
     52         if self.fo:
     53             self.fo.close()
     54         if self.name and os.path.exists(self.name):
     55             os.remove(self.name)
     56 
     57         self.fd = self.fo = self.name = None
     58 
     59 
     60     def __del__(self):
     61         try:
     62             if self.name is not None:
     63                 logging.debug('Clean was not called for ' + self.name)
     64                 self.clean()
     65         except:
     66             try:
     67                 msg = 'An exception occurred while calling the destructor'
     68                 logging.exception(msg)
     69             except:
     70                 pass
     71 
     72 
     73 class tempdir(object):
     74     """
     75     A wrapper for tempfile.mkdtemp
     76 
     77     @var name: The name of the temporary dir.
     78     @return A tempdir object
     79     example usage:
     80         b = autotemp.tempdir(unique_id='exemdir')
     81         b.name # your directory
     82         b.clean() # clean up after yourself
     83     """
     84     def __init__(self,  suffix='', unique_id='', prefix='', dir=None):
     85         """
     86         Initialize temp directory.
     87 
     88         @param suffix: suffix for dir.
     89         @param prefix: prefix for dir. Defaults to '_autotmp'.
     90         @param unique_id: unique id of tempdir.
     91         @param dir: parent directory of the tempdir. Defaults to /tmp.
     92 
     93         eg: autotemp.tempdir(suffix='suffix', unique_id='123', prefix='prefix')
     94             creates a dir like '/tmp/prefix_autotmp_<random hash>123suffix'
     95         """
     96         suffix = unique_id + suffix
     97         prefix = prefix + _TEMPLATE
     98         self.name = module_tempfile.mkdtemp(suffix=suffix,
     99                                             prefix=prefix, dir=dir)
    100 
    101 
    102     def clean(self):
    103         """
    104         Remove the temporary dir that was created.
    105         This is also called by the destructor.
    106         """
    107         if self.name and os.path.exists(self.name):
    108             shutil.rmtree(self.name)
    109 
    110         self.name = None
    111 
    112 
    113     def __del__(self):
    114         try:
    115             if self.name:
    116                 logging.debug('Clean was not called for ' + self.name)
    117                 self.clean()
    118         except:
    119             try:
    120                 msg = 'An exception occurred while calling the destructor'
    121                 logging.exception(msg)
    122             except:
    123                 pass
    124 
    125 
    126 class dummy_dir(object):
    127     """A dummy object representing a directory with a name.
    128 
    129     Only used for compat with the tmpdir, in cases where we wish to
    130     reuse a dir with the same interface but not to delete it after
    131     we're done using it.
    132     """
    133 
    134     def __init__(self, name):
    135         """Initialize the dummy_dir object.
    136 
    137         @param name: Path the the directory.
    138         """
    139         self.name = name
    140