• 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('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