• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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('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
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        """
86        Initialize temp directory.
87
88        @param suffix: suffix for dir.
89        @param prefix: prefix for dir. Defaults to '_autotmp'.
90        @param unique_id: unique id of tempdir.
91        @param dir: parent directory of the tempdir. Defaults to /tmp.
92
93        eg: autotemp.tempdir(suffix='suffix', unique_id='123', prefix='prefix')
94            creates a dir like '/tmp/prefix_autotmp_<random hash>123suffix'
95        """
96        suffix = unique_id + suffix
97        prefix = prefix + _TEMPLATE
98        self.name = module_tempfile.mkdtemp(suffix=suffix,
99                                            prefix=prefix, dir=dir)
100
101
102    def clean(self):
103        """
104        Remove the temporary dir that was created.
105        This is also called by the destructor.
106        """
107        if self.name and os.path.exists(self.name):
108            shutil.rmtree(self.name)
109
110        self.name = None
111
112
113    def __del__(self):
114        try:
115            if self.name:
116                logging.debug('Clean was not called for ' + self.name)
117                self.clean()
118        except:
119            try:
120                msg = 'An exception occurred while calling the destructor'
121                logging.exception(msg)
122            except:
123                pass
124
125
126class dummy_dir(object):
127    """A dummy object representing a directory with a name.
128
129    Only used for compat with the tmpdir, in cases where we wish to
130    reuse a dir with the same interface but not to delete it after
131    we're done using it.
132    """
133
134    def __init__(self, name):
135        """Initialize the dummy_dir object.
136
137        @param name: Path the the directory.
138        """
139        self.name = name
140