mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 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 | |
jadmanski | 5f2bc28 | 2009-07-28 17:21:15 +0000 | [diff] [blame] | 15 | import shutil, os, logging |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 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 |
Eric Li | e0493a4 | 2010-11-15 13:05:43 -0800 | [diff] [blame] | 34 | t.fo # file object |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 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) |
Eric Li | e0493a4 | 2010-11-15 13:05:43 -0800 | [diff] [blame] | 44 | self.fo = os.fdopen(self.fd) |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 45 | |
| 46 | |
| 47 | def clean(self): |
| 48 | """ |
| 49 | Remove the temporary file that was created. |
| 50 | This is also called by the destructor. |
| 51 | """ |
Eric Li | e0493a4 | 2010-11-15 13:05:43 -0800 | [diff] [blame] | 52 | if self.fo: |
| 53 | self.fo.close() |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 54 | if self.name and os.path.exists(self.name): |
| 55 | os.remove(self.name) |
| 56 | |
Eric Li | e0493a4 | 2010-11-15 13:05:43 -0800 | [diff] [blame] | 57 | self.fd = self.fo = self.name = None |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 58 | |
| 59 | |
| 60 | def __del__(self): |
| 61 | try: |
Eric Li | e0493a4 | 2010-11-15 13:05:43 -0800 | [diff] [blame] | 62 | if self.name is not None: |
Congbin Guo | 66abb66 | 2018-02-12 16:55:24 -0800 | [diff] [blame] | 63 | logging.debug('Cleaning %s', self.name) |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 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: |
jadmanski | 5f2bc28 | 2009-07-28 17:21:15 +0000 | [diff] [blame] | 80 | b = autotemp.tempdir(unique_id='exemdir') |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 81 | b.name # your directory |
| 82 | b.clean() # clean up after yourself |
| 83 | """ |
Ningning Xia | 84190b8 | 2018-04-16 15:01:40 -0700 | [diff] [blame] | 84 | def __init__(self, suffix='', unique_id='', prefix='', dir=None, |
| 85 | auto_clean=True): |
beeps | 98365d8 | 2013-02-20 20:08:07 -0800 | [diff] [blame] | 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. |
Ningning Xia | 84190b8 | 2018-04-16 15:01:40 -0700 | [diff] [blame] | 93 | @param auto_clean: automatically clean up the tempdir in destructor. |
beeps | 98365d8 | 2013-02-20 20:08:07 -0800 | [diff] [blame] | 94 | |
| 95 | eg: autotemp.tempdir(suffix='suffix', unique_id='123', prefix='prefix') |
| 96 | creates a dir like '/tmp/prefix_autotmp_<random hash>123suffix' |
| 97 | """ |
Ningning Xia | 84190b8 | 2018-04-16 15:01:40 -0700 | [diff] [blame] | 98 | self.auto_clean = auto_clean |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 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: |
Ningning Xia | 84190b8 | 2018-04-16 15:01:40 -0700 | [diff] [blame] | 118 | if self.name and self.auto_clean: |
mbligh | 6a242df | 2009-07-02 19:00:02 +0000 | [diff] [blame] | 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 |
Prashanth Balasubramanian | f571aa6 | 2014-10-13 18:09:44 -0700 | [diff] [blame] | 127 | |
| 128 | |
| 129 | class 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 |