1""" 2Autotest tempfile wrapper for mkstemp (known as tempfile here) and 3mkdtemp (known as tempdir). 4 5This wrapper provides a mechanism to clean up temporary files/dirs once they 6are no longer need. 7 8Files/Dirs will have a unique_id prepended to the suffix and a 9_autotmp_ tag appended to the prefix. 10 11It is required that the unique_id param is supplied when a temp dir/file is 12created. 13""" 14 15import shutil, os, logging 16import tempfile as module_tempfile 17 18_TEMPLATE = '_autotmp_' 19 20 21class 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('Cleaning %s', 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 73class 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 auto_clean=True): 86 """ 87 Initialize temp directory. 88 89 @param suffix: suffix for dir. 90 @param prefix: prefix for dir. Defaults to '_autotmp'. 91 @param unique_id: unique id of tempdir. 92 @param dir: parent directory of the tempdir. Defaults to /tmp. 93 @param auto_clean: automatically clean up the tempdir in destructor. 94 95 eg: autotemp.tempdir(suffix='suffix', unique_id='123', prefix='prefix') 96 creates a dir like '/tmp/prefix_autotmp_<random hash>123suffix' 97 """ 98 self.auto_clean = auto_clean 99 suffix = unique_id + suffix 100 prefix = prefix + _TEMPLATE 101 self.name = module_tempfile.mkdtemp(suffix=suffix, 102 prefix=prefix, dir=dir) 103 104 105 def clean(self): 106 """ 107 Remove the temporary dir that was created. 108 This is also called by the destructor. 109 """ 110 if self.name and os.path.exists(self.name): 111 shutil.rmtree(self.name) 112 113 self.name = None 114 115 116 def __del__(self): 117 try: 118 if self.name and self.auto_clean: 119 logging.debug('Clean was not called for ' + self.name) 120 self.clean() 121 except: 122 try: 123 msg = 'An exception occurred while calling the destructor' 124 logging.exception(msg) 125 except: 126 pass 127 128 129class dummy_dir(object): 130 """A dummy object representing a directory with a name. 131 132 Only used for compat with the tmpdir, in cases where we wish to 133 reuse a dir with the same interface but not to delete it after 134 we're done using it. 135 """ 136 137 def __init__(self, name): 138 """Initialize the dummy_dir object. 139 140 @param name: Path the the directory. 141 """ 142 self.name = name 143