PageRenderTime 97ms CodeModel.GetById 65ms app.highlight 27ms RepoModel.GetById 1ms app.codeStats 0ms

/Lib/tempfile.py

http://unladen-swallow.googlecode.com/
Python | 603 lines | 599 code | 0 blank | 4 comment | 1 complexity | 1eca671f8c81e26d7a78da7c62270c5d MD5 | raw file
  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
415
416def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",
417                       prefix=template, dir=None, delete=True):
418    """Create and return a temporary file.
419    Arguments:
420    'prefix', 'suffix', 'dir' -- as for mkstemp.
421    'mode' -- the mode argument to os.fdopen (default "w+b").
422    'bufsize' -- the buffer size argument to os.fdopen (default -1).
423    'delete' -- whether the file is deleted on close (default True).
424    The file is created as mkstemp() would do it.
425
426    Returns an object with a file-like interface; the name of the file
427    is accessible as file.name.  The file will be automatically deleted
428    when it is closed unless the 'delete' argument is set to False.
429    """
430
431    if dir is None:
432        dir = gettempdir()
433
434    if 'b' in mode:
435        flags = _bin_openflags
436    else:
437        flags = _text_openflags
438
439    # Setting O_TEMPORARY in the flags causes the OS to delete
440    # the file when it is closed.  This is only supported by Windows.
441    if _os.name == 'nt' and delete:
442        flags |= _os.O_TEMPORARY
443
444    (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
445    file = _os.fdopen(fd, mode, bufsize)
446    return _TemporaryFileWrapper(file, name, delete)
447
448if _os.name != 'posix' or _os.sys.platform == 'cygwin':
449    # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
450    # while it is open.
451    TemporaryFile = NamedTemporaryFile
452
453else:
454    def TemporaryFile(mode='w+b', bufsize=-1, suffix="",
455                      prefix=template, dir=None):
456        """Create and return a temporary file.
457        Arguments:
458        'prefix', 'suffix', 'dir' -- as for mkstemp.
459        'mode' -- the mode argument to os.fdopen (default "w+b").
460        'bufsize' -- the buffer size argument to os.fdopen (default -1).
461        The file is created as mkstemp() would do it.
462
463        Returns an object with a file-like interface.  The file has no
464        name, and will cease to exist when it is closed.
465        """
466
467        if dir is None:
468            dir = gettempdir()
469
470        if 'b' in mode:
471            flags = _bin_openflags
472        else:
473            flags = _text_openflags
474
475        (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
476        try:
477            _os.unlink(name)
478            return _os.fdopen(fd, mode, bufsize)
479        except:
480            _os.close(fd)
481            raise
482
483class SpooledTemporaryFile:
484    """Temporary file wrapper, specialized to switch from
485    StringIO to a real file when it exceeds a certain size or
486    when a fileno is needed.
487    """
488    _rolled = False
489
490    def __init__(self, max_size=0, mode='w+b', bufsize=-1,
491                 suffix="", prefix=template, dir=None):
492        self._file = _StringIO()
493        self._max_size = max_size
494        self._rolled = False
495        self._TemporaryFileArgs = (mode, bufsize, suffix, prefix, dir)
496
497    def _check(self, file):
498        if self._rolled: return
499        max_size = self._max_size
500        if max_size and file.tell() > max_size:
501            self.rollover()
502
503    def rollover(self):
504        if self._rolled: return
505        file = self._file
506        newfile = self._file = TemporaryFile(*self._TemporaryFileArgs)
507        del self._TemporaryFileArgs
508
509        newfile.write(file.getvalue())
510        newfile.seek(file.tell(), 0)
511
512        self._rolled = True
513
514    # The method caching trick from NamedTemporaryFile
515    # won't work here, because _file may change from a
516    # _StringIO instance to a real file. So we list
517    # all the methods directly.
518
519    # Context management protocol
520    def __enter__(self):
521        if self._file.closed:
522            raise ValueError("Cannot enter context with closed file")
523        return self
524
525    def __exit__(self, exc, value, tb):
526        self._file.close()
527
528    # file protocol
529    def __iter__(self):
530        return self._file.__iter__()
531
532    def close(self):
533        self._file.close()
534
535    @property
536    def closed(self):
537        return self._file.closed
538
539    @property
540    def encoding(self):
541        return self._file.encoding
542
543    def fileno(self):
544        self.rollover()
545        return self._file.fileno()
546
547    def flush(self):
548        self._file.flush()
549
550    def isatty(self):
551        return self._file.isatty()
552
553    @property
554    def mode(self):
555        return self._file.mode
556
557    @property
558    def name(self):
559        return self._file.name
560
561    @property
562    def newlines(self):
563        return self._file.newlines
564
565    def next(self):
566        return self._file.next
567
568    def read(self, *args):
569        return self._file.read(*args)
570
571    def readline(self, *args):
572        return self._file.readline(*args)
573
574    def readlines(self, *args):
575        return self._file.readlines(*args)
576
577    def seek(self, *args):
578        self._file.seek(*args)
579
580    @property
581    def softspace(self):
582        return self._file.softspace
583
584    def tell(self):
585        return self._file.tell()
586
587    def truncate(self):
588        self._file.truncate()
589
590    def write(self, s):
591        file = self._file
592        rv = file.write(s)
593        self._check(file)
594        return rv
595
596    def writelines(self, iterable):
597        file = self._file
598        rv = file.writelines(iterable)
599        self._check(file)
600        return rv
601
602    def xreadlines(self, *args):
603        return self._file.xreadlines(*args)