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