client/common_lib/autotemp.py

   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=None, prefix='', dir=None):
  85         suffix = unique_id + suffix
  86         prefix = prefix + _TEMPLATE
  87         self.name = module_tempfile.mkdtemp(suffix=suffix,
  88                                             prefix=prefix, dir=dir)
  89
  90
  91     def clean(self):
  92         """
  93         Remove the temporary dir that was created.
  94         This is also called by the destructor.
  95         """
  96         if self.name and os.path.exists(self.name):
  97             shutil.rmtree(self.name)
  98
  99         self.name = None
 100
 101
 102     def __del__(self):
 103         try:
 104             if self.name:
 105                 logging.debug('Clean was not called for ' + self.name)
 106                 self.clean()
 107         except:
 108             try:
 109                 msg = 'An exception occurred while calling the destructor'
 110                 logging.exception(msg)
 111             except:
 112                 pass