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