1"""Temporary files.
2
3This module provides generic, low- and high-level interfaces for
4creating temporary files and directories.  The interfaces listed
5as "safe" just below can be used without fear of race conditions.
6Those listed as "unsafe" cannot, and are provided for backward
7compatibility only.
8
9This module also provides some data items to the user:
10
11  TMP_MAX  - maximum number of names that will be tried before
12             giving up.
13  template - the default prefix for all temporary names.
14             You may change this to control the default prefix.
15  tempdir  - If this is set to a string before the first use of
16             any routine from this module, it will be considered as
17             another candidate location to store temporary files.
18"""
19
20__all__ = [
21    "NamedTemporaryFile", "TemporaryFile", # high level safe interfaces
22    "SpooledTemporaryFile",
23    "mkstemp", "mkdtemp",                  # low level safe interfaces
24    "mktemp",                              # deprecated unsafe interface
25    "TMP_MAX", "gettempprefix",            # constants
26    "tempdir", "gettempdir"
27   ]
28
29
30# Imports.
31
32import os as _os
33import errno as _errno
34from random import Random as _Random
35
36try:
37    from cStringIO import StringIO as _StringIO
38except ImportError:
39    from StringIO import StringIO as _StringIO
40
41try:
42    import fcntl as _fcntl
43except ImportError:
44    def _set_cloexec(fd):
45        pass
46else:
47    def _set_cloexec(fd):
48        try:
49            flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)
50        except IOError:
51            pass
52        else:
53            # flags read successfully, modify
54            flags |= _fcntl.FD_CLOEXEC
55            _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)
56
57
58try:
59    import thread as _thread
60except ImportError:
61    import dummy_thread as _thread
62_allocate_lock = _thread.allocate_lock
63
64_text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL
65if hasattr(_os, 'O_NOINHERIT'):
66    _text_openflags |= _os.O_NOINHERIT
67if hasattr(_os, 'O_NOFOLLOW'):
68    _text_openflags |= _os.O_NOFOLLOW
69
70_bin_openflags = _text_openflags
71if hasattr(_os, 'O_BINARY'):
72    _bin_openflags |= _os.O_BINARY
73
74if hasattr(_os, 'TMP_MAX'):
75    TMP_MAX = _os.TMP_MAX
76else:
77    TMP_MAX = 10000
78
79template = "tmp"
80
81# Internal routines.
82
83_once_lock = _allocate_lock()
84
85if hasattr(_os, "lstat"):
86    _stat = _os.lstat
87elif hasattr(_os, "stat"):
88    _stat = _os.stat
89else:
90    # Fallback.  All we need is something that raises os.error if the
91    # file doesn't exist.
92    def _stat(fn):
93        try:
94            f = open(fn)
95        except IOError:
96            raise _os.error
97        f.close()
98
99def _exists(fn):
100    try:
101        _stat(fn)
102    except _os.error:
103        return False
104    else:
105        return True
106
107class _RandomNameSequence:
108    """An instance of _RandomNameSequence generates an endless
109    sequence of unpredictable strings which can safely be incorporated
110    into file names.  Each string is six characters long.  Multiple
111    threads can safely use the same instance at the same time.
112
113    _RandomNameSequence is an iterator."""
114
115    characters = ("abcdefghijklmnopqrstuvwxyz" +
116                  "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
117                  "0123456789_")
118
119    def __init__(self):
120        self.mutex = _allocate_lock()
121        self.rng = _Random()
122        self.normcase = _os.path.normcase
123
124    def __iter__(self):
125        return self
126
127    def next(self):
128        m = self.mutex
129        c = self.characters
130        choose = self.rng.choice
131
132        m.acquire()
133        try:
134            letters = [choose(c) for dummy in "123456"]
135        finally:
136            m.release()
137
138        return self.normcase(''.join(letters))
139
140def _candidate_tempdir_list():
141    """Generate a list of candidate temporary directories which
142    _get_default_tempdir will try."""
143
144    dirlist = []
145
146    # First, try the environment.
147    for envname in 'TMPDIR', 'TEMP', 'TMP':
148        dirname = _os.getenv(envname)
149        if dirname: dirlist.append(dirname)
150
151    # Failing that, try OS-specific locations.
152    if _os.name == 'riscos':
153        dirname = _os.getenv('Wimp$ScrapDir')
154        if dirname: dirlist.append(dirname)
155    elif _os.name == 'nt':
156        dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
157    else:
158        dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])
159
160    # As a last resort, the current directory.
161    try:
162        dirlist.append(_os.getcwd())
163    except (AttributeError, _os.error):
164        dirlist.append(_os.curdir)
165
166    return dirlist
167
168def _get_default_tempdir():
169    """Calculate the default directory to use for temporary files.
170    This routine should be called exactly once.
171
172    We determine whether or not a candidate temp dir is usable by
173    trying to create and write to a file in that directory.  If this
174    is successful, the test file is deleted.  To prevent denial of
175    service, the name of the test file must be randomized."""
176
177    namer = _RandomNameSequence()
178    dirlist = _candidate_tempdir_list()
179    flags = _text_openflags
180
181    for dir in dirlist:
182        if dir != _os.curdir:
183            dir = _os.path.normcase(_os.path.abspath(dir))
184        # Try only a few names per directory.
185        for seq in xrange(100):
186            name = namer.next()
187            filename = _os.path.join(dir, name)
188            try:
189                fd = _os.open(filename, flags, 0600)
190                fp = _os.fdopen(fd, 'w')
191                fp.write('blat')
192                fp.close()
193                _os.unlink(filename)
194                del fp, fd
195                return dir
196            except (OSError, IOError), e:
197                if e[0] != _errno.EEXIST:
198                    break # no point trying more names in this directory
199                pass
200    raise IOError, (_errno.ENOENT,
201                    ("No usable temporary directory found in %s" % dirlist))
202
203_name_sequence = None
204
205def _get_candidate_names():
206    """Common setup sequence for all user-callable interfaces."""
207
208    global _name_sequence
209    if _name_sequence is None:
210        _once_lock.acquire()
211        try:
212            if _name_sequence is None:
213                _name_sequence = _RandomNameSequence()
214        finally:
215            _once_lock.release()
216    return _name_sequence
217
218
219def _mkstemp_inner(dir, pre, suf, flags):
220    """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
221
222    names = _get_candidate_names()
223
224    for seq in xrange(TMP_MAX):
225        name = names.next()
226        file = _os.path.join(dir, pre + name + suf)
227        try:
228            fd = _os.open(file, flags, 0600)
229            _set_cloexec(fd)
230            return (fd, _os.path.abspath(file))
231        except OSError, e:
232            if e.errno == _errno.EEXIST:
233                continue # try again
234            raise
235
236    raise IOError, (_errno.EEXIST, "No usable temporary file name found")
237
238
239# User visible interfaces.
240
241def gettempprefix():
242    """Accessor for tempdir.template."""
243    return template
244
245tempdir = None
246
247def gettempdir():
248    """Accessor for tempfile.tempdir."""
249    global tempdir
250    if tempdir is None:
251        _once_lock.acquire()
252        try:
253            if tempdir is None:
254                tempdir = _get_default_tempdir()
255        finally:
256            _once_lock.release()
257    return tempdir
258
259def mkstemp(suffix="", prefix=template, dir=None, text=False):
260    """User-callable function to create and return a unique temporary
261    file.  The return value is a pair (fd, name) where fd is the
262    file descriptor returned by os.open, and name is the filename.
263
264    If 'suffix' is specified, the file name will end with that suffix,
265    otherwise there will be no suffix.
266
267    If 'prefix' is specified, the file name will begin with that prefix,
268    otherwise a default prefix is used.
269
270    If 'dir' is specified, the file will be created in that directory,
271    otherwise a default directory is used.
272
273    If 'text' is specified and true, the file is opened in text
274    mode.  Else (the default) the file is opened in binary mode.  On
275    some operating systems, this makes no difference.
276
277    The file is readable and writable only by the creating user ID.
278    If the operating system uses permission bits to indicate whether a
279    file is executable, the file is executable by no one. The file
280    descriptor is not inherited by children of this process.
281
282    Caller is responsible for deleting the file when done with it.
283    """
284
285    if dir is None:
286        dir = gettempdir()
287
288    if text:
289        flags = _text_openflags
290    else:
291        flags = _bin_openflags
292
293    return _mkstemp_inner(dir, prefix, suffix, flags)
294
295
296def mkdtemp(suffix="", prefix=template, dir=None):
297    """User-callable function to create and return a unique temporary
298    directory.  The return value is the pathname of the directory.
299
300    Arguments are as for mkstemp, except that the 'text' argument is
301    not accepted.
302
303    The directory is readable, writable, and searchable only by the
304    creating user.
305
306    Caller is responsible for deleting the directory when done with it.
307    """
308
309    if dir is None:
310        dir = gettempdir()
311
312    names = _get_candidate_names()
313
314    for seq in xrange(TMP_MAX):
315        name = names.next()
316        file = _os.path.join(dir, prefix + name + suffix)
317        try:
318            _os.mkdir(file, 0700)
319            return file
320        except OSError, e:
321            if e.errno == _errno.EEXIST:
322                continue # try again
323            raise
324
325    raise IOError, (_errno.EEXIST, "No usable temporary directory name found")
326
327def mktemp(suffix="", prefix=template, dir=None):
328    """User-callable function to return a unique temporary file name.  The
329    file is not created.
330
331    Arguments are as for mkstemp, except that the 'text' argument is
332    not accepted.
333
334    This function is unsafe and should not be used.  The file name
335    refers to a file that did not exist at some point, but by the time
336    you get around to creating it, someone else may have beaten you to
337    the punch.
338    """
339
340##    from warnings import warn as _warn
341##    _warn("mktemp is a potential security risk to your program",
342##          RuntimeWarning, stacklevel=2)
343
344    if dir is None:
345        dir = gettempdir()
346
347    names = _get_candidate_names()
348    for seq in xrange(TMP_MAX):
349        name = names.next()
350        file = _os.path.join(dir, prefix + name + suffix)
351        if not _exists(file):
352            return file
353
354    raise IOError, (_errno.EEXIST, "No usable temporary filename found")
355
356
357class _TemporaryFileWrapper:
358    """Temporary file wrapper
359
360    This class provides a wrapper around files opened for
361    temporary use.  In particular, it seeks to automatically
362    remove the file when it is no longer needed.
363    """
364
365    def __init__(self, file, name, delete=True):
366        self.file = file
367        self.name = name
368        self.close_called = False
369        self.delete = delete
370
371    def __getattr__(self, name):
372        # Attribute lookups are delegated to the underlying file
373        # and cached for non-numeric results
374        # (i.e. methods are cached, closed and friends are not)
375        file = self.__dict__['file']
376        a = getattr(file, name)
377        if not issubclass(type(a), type(0)):
378            setattr(self, name, a)
379        return a
380
381    # The underlying __enter__ method returns the wrong object
382    # (self.file) so override it to return the wrapper
383    def __enter__(self):
384        self.file.__enter__()
385        return self
386
387    # NT provides delete-on-close as a primitive, so we don't need
388    # the wrapper to do anything special.  We still use it so that
389    # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
390    if _os.name != 'nt':
391        # Cache the unlinker so we don't get spurious errors at
392        # shutdown when the module-level "os" is None'd out.  Note
393        # that this must be referenced as self.unlink, because the
394        # name TemporaryFileWrapper may also get None'd out before
395        # __del__ is called.
396        unlink = _os.unlink
397
398        def close(self):
399            if not self.close_called:
400                self.close_called = True
401                self.file.close()
402                if self.delete:
403                    self.unlink(self.name)
404
405        def __del__(self):
406            self.close()
407
408        # Need to trap __exit__ as well to ensure the file gets
409        # deleted when used in a with statement
410        def __exit__(self, exc, value, tb):
411            result = self.file.__exit__(exc, value, tb)
412            self.close()
413            return result
414    else:
415        def __exit__(self, exc, value, tb):
416            self.file.__exit__(exc, value, tb)
417
418
419def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",
420                       prefix=template, dir=None, delete=True):
421    """Create and return a temporary file.
422    Arguments:
423    'prefix', 'suffix', 'dir' -- as for mkstemp.
424    'mode' -- the mode argument to os.fdopen (default "w+b").
425    'bufsize' -- the buffer size argument to os.fdopen (default -1).
426    'delete' -- whether the file is deleted on close (default True).
427    The file is created as mkstemp() would do it.
428
429    Returns an object with a file-like interface; the name of the file
430    is accessible as file.name.  The file will be automatically deleted
431    when it is closed unless the 'delete' argument is set to False.
432    """
433
434    if dir is None:
435        dir = gettempdir()
436
437    if 'b' in mode:
438        flags = _bin_openflags
439    else:
440        flags = _text_openflags
441
442    # Setting O_TEMPORARY in the flags causes the OS to delete
443    # the file when it is closed.  This is only supported by Windows.
444    if _os.name == 'nt' and delete:
445        flags |= _os.O_TEMPORARY
446
447    (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
448    file = _os.fdopen(fd, mode, bufsize)
449    return _TemporaryFileWrapper(file, name, delete)
450
451if _os.name != 'posix' or _os.sys.platform == 'cygwin':
452    # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
453    # while it is open.
454    TemporaryFile = NamedTemporaryFile
455
456else:
457    def TemporaryFile(mode='w+b', bufsize=-1, suffix="",
458                      prefix=template, dir=None):
459        """Create and return a temporary file.
460        Arguments:
461        'prefix', 'suffix', 'dir' -- as for mkstemp.
462        'mode' -- the mode argument to os.fdopen (default "w+b").
463        'bufsize' -- the buffer size argument to os.fdopen (default -1).
464        The file is created as mkstemp() would do it.
465
466        Returns an object with a file-like interface.  The file has no
467        name, and will cease to exist when it is closed.
468        """
469
470        if dir is None:
471            dir = gettempdir()
472
473        if 'b' in mode:
474            flags = _bin_openflags
475        else:
476            flags = _text_openflags
477
478        (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
479        try:
480            _os.unlink(name)
481            return _os.fdopen(fd, mode, bufsize)
482        except:
483            _os.close(fd)
484            raise
485
486class SpooledTemporaryFile:
487    """Temporary file wrapper, specialized to switch from
488    StringIO to a real file when it exceeds a certain size or
489    when a fileno is needed.
490    """
491    _rolled = False
492
493    def __init__(self, max_size=0, mode='w+b', bufsize=-1,
494                 suffix="", prefix=template, dir=None):
495        self._file = _StringIO()
496        self._max_size = max_size
497        self._rolled = False
498        self._TemporaryFileArgs = (mode, bufsize, suffix, prefix, dir)
499
500    def _check(self, file):
501        if self._rolled: return
502        max_size = self._max_size
503        if max_size and file.tell() > max_size:
504            self.rollover()
505
506    def rollover(self):
507        if self._rolled: return
508        file = self._file
509        newfile = self._file = TemporaryFile(*self._TemporaryFileArgs)
510        del self._TemporaryFileArgs
511
512        newfile.write(file.getvalue())
513        newfile.seek(file.tell(), 0)
514
515        self._rolled = True
516
517    # The method caching trick from NamedTemporaryFile
518    # won't work here, because _file may change from a
519    # _StringIO instance to a real file. So we list
520    # all the methods directly.
521
522    # Context management protocol
523    def __enter__(self):
524        if self._file.closed:
525            raise ValueError("Cannot enter context with closed file")
526        return self
527
528    def __exit__(self, exc, value, tb):
529        self._file.close()
530
531    # file protocol
532    def __iter__(self):
533        return self._file.__iter__()
534
535    def close(self):
536        self._file.close()
537
538    @property
539    def closed(self):
540        return self._file.closed
541
542    @property
543    def encoding(self):
544        return self._file.encoding
545
546    def fileno(self):
547        self.rollover()
548        return self._file.fileno()
549
550    def flush(self):
551        self._file.flush()
552
553    def isatty(self):
554        return self._file.isatty()
555
556    @property
557    def mode(self):
558        return self._file.mode
559
560    @property
561    def name(self):
562        return self._file.name
563
564    @property
565    def newlines(self):
566        return self._file.newlines
567
568    def next(self):
569        return self._file.next
570
571    def read(self, *args):
572        return self._file.read(*args)
573
574    def readline(self, *args):
575        return self._file.readline(*args)
576
577    def readlines(self, *args):
578        return self._file.readlines(*args)
579
580    def seek(self, *args):
581        self._file.seek(*args)
582
583    @property
584    def softspace(self):
585        return self._file.softspace
586
587    def tell(self):
588        return self._file.tell()
589
590    def truncate(self):
591        self._file.truncate()
592
593    def write(self, s):
594        file = self._file
595        rv = file.write(s)
596        self._check(file)
597        return rv
598
599    def writelines(self, iterable):
600        file = self._file
601        rv = file.writelines(iterable)
602        self._check(file)
603        return rv
604
605    def xreadlines(self, *args):
606        return self._file.xreadlines(*args)
607