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 @param auto_clean: automatically delete the temporary file in destructor. 28 @var name: The name of the temporary file. 29 @var fd: the file descriptor of the temporary file that was created. 30 @return a tempfile object 31 example usage: 32 t = autotemp.tempfile(unique_id='fig') 33 t.name # name of file 34 t.fd # file descriptor 35 t.fo # file object 36 t.clean() # clean up after yourself 37 """ 38 def __init__(self, unique_id, suffix='', prefix='', dir=None, 39 text=False, auto_clean=True): 40 self.auto_clean = auto_clean 41 suffix = unique_id + suffix 42 prefix = prefix + _TEMPLATE 43 self.fd, self.name = module_tempfile.mkstemp(suffix=suffix, 44 prefix=prefix, 45 dir=dir, text=text) 46 self.fo = os.fdopen(self.fd) 47 48 49 def clean(self): 50 """ 51 Remove the temporary file that was created. 52 This is also called by the destructor. 53 """ 54 if self.fo: 55 self.fo.close() 56 if self.name and os.path.exists(self.name): 57 os.remove(self.name) 58 59 self.fd = self.fo = self.name = None 60 61 62 def __del__(self): 63 try: 64 if self.name is not None and self.auto_clean: 65 logging.debug('Auto-cleaning %s', self.name) 66 self.clean() 67 except: 68 try: 69 msg = 'An exception occurred while calling the destructor' 70 logging.exception(msg) 71 except: 72 pass 73 74 75class tempdir(object): 76 """ 77 A wrapper for tempfile.mkdtemp 78 79 @var name: The name of the temporary dir. 80 @return A tempdir object 81 example usage: 82 b = autotemp.tempdir(unique_id='exemdir') 83 b.name # your directory 84 b.clean() # clean up after yourself 85 """ 86 def __init__(self, suffix='', unique_id='', prefix='', dir=None, 87 auto_clean=True): 88 """ 89 Initialize temp directory. 90 91 @param suffix: suffix for dir. 92 @param prefix: prefix for dir. Defaults to '_autotmp'. 93 @param unique_id: unique id of tempdir. 94 @param dir: parent directory of the tempdir. Defaults to /tmp. 95 @param auto_clean: automatically clean up the tempdir in destructor. 96 97 eg: autotemp.tempdir(suffix='suffix', unique_id='123', prefix='prefix') 98 creates a dir like '/tmp/prefix_autotmp_<random hash>123suffix' 99 """ 100 self.auto_clean = auto_clean 101 suffix = unique_id + suffix 102 prefix = prefix + _TEMPLATE 103 self.name = module_tempfile.mkdtemp(suffix=suffix, 104 prefix=prefix, dir=dir) 105 106 107 def clean(self): 108 """ 109 Remove the temporary dir that was created. 110 This is also called by the destructor. 111 """ 112 if self.name and os.path.exists(self.name): 113 shutil.rmtree(self.name) 114 115 self.name = None 116 117 118 def __del__(self): 119 try: 120 if self.name and self.auto_clean: 121 logging.debug('Auto-cleaning %s', self.name) 122 self.clean() 123 except: 124 try: 125 msg = 'An exception occurred while calling the destructor' 126 logging.exception(msg) 127 except: 128 pass 129 130 131class dummy_dir(object): 132 """A dummy object representing a directory with a name. 133 134 Only used for compat with the tmpdir, in cases where we wish to 135 reuse a dir with the same interface but not to delete it after 136 we're done using it. 137 """ 138 139 def __init__(self, name): 140 """Initialize the dummy_dir object. 141 142 @param name: Path the the directory. 143 """ 144 self.name = name 145