tempfile.py 13.6 KB
Newer Older
1
"""Temporary files.
2

3 4 5 6 7
This module provides generic, low- and high-level interfaces for
creating temporary files and directories.  The interfaces listed
as "safe" just below can be used without fear of race conditions.
Those listed as "unsafe" cannot, and are provided for backward
compatibility only.
Guido van Rossum's avatar
Guido van Rossum committed
8

9
This module also provides some data items to the user:
Guido van Rossum's avatar
Guido van Rossum committed
10

11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
  TMP_MAX  - maximum number of names that will be tried before
             giving up.
  template - the default prefix for all temporary names.
             You may change this to control the default prefix.
  tempdir  - If this is set to a string before the first use of
             any routine from this module, it will be considered as
             another candidate location to store temporary files.
"""

__all__ = [
    "NamedTemporaryFile", "TemporaryFile", # high level safe interfaces
    "mkstemp", "mkdtemp",                  # low level safe interfaces
    "mktemp",                              # deprecated unsafe interface
    "TMP_MAX", "gettempprefix",            # constants
    "tempdir", "gettempdir"
   ]


# Imports.

import os as _os
import errno as _errno
from random import Random as _Random

if _os.name == 'mac':
    import macfs as _macfs
    import MACFS as _MACFS

try:
    import fcntl as _fcntl
    def _set_cloexec(fd):
        flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)
        if flags >= 0:
            # flags read successfully, modify
            flags |= _fcntl.FD_CLOEXEC
            _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)
except (ImportError, AttributeError):
    def _set_cloexec(fd):
        pass

try:
    import thread as _thread
    _allocate_lock = _thread.allocate_lock
except (ImportError, AttributeError):
    class _allocate_lock:
        def acquire(self):
            pass
        release = acquire

_text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL
Tim Peters's avatar
Tim Peters committed
61 62 63 64
if hasattr(_os, 'O_NOINHERIT'):
    _text_openflags |= _os.O_NOINHERIT
if hasattr(_os, 'O_NOFOLLOW'):
    _text_openflags |= _os.O_NOFOLLOW
65 66

_bin_openflags = _text_openflags
Tim Peters's avatar
Tim Peters committed
67 68
if hasattr(_os, 'O_BINARY'):
    _bin_openflags |= _os.O_BINARY
69 70 71 72 73 74

if hasattr(_os, 'TMP_MAX'):
    TMP_MAX = _os.TMP_MAX
else:
    TMP_MAX = 10000

75
template = "tmp"
76

77
tempdir = None
78

79 80 81 82 83 84 85 86 87 88 89 90
# Internal routines.

_once_lock = _allocate_lock()

class _RandomNameSequence:
    """An instance of _RandomNameSequence generates an endless
    sequence of unpredictable strings which can safely be incorporated
    into file names.  Each string is six characters long.  Multiple
    threads can safely use the same instance at the same time.

    _RandomNameSequence is an iterator."""

91 92 93
    characters = ("abcdefghijklmnopqrstuvwxyz" +
                  "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
                  "0123456789-_")
94 95 96 97 98

    def __init__(self):
        self.mutex = _allocate_lock()
        self.rng = _Random()
        self.normcase = _os.path.normcase
99

100 101 102 103 104 105
    def __iter__(self):
        return self

    def next(self):
        m = self.mutex
        c = self.characters
106
        choose = self.rng.choice
107

108
        m.acquire()
109
        try:
110
            letters = [choose(c) for dummy in "123456"]
111 112 113
        finally:
            m.release()

114
        return self.normcase(''.join(letters))
115 116 117 118 119 120 121 122

def _candidate_tempdir_list():
    """Generate a list of candidate temporary directories which
    _get_default_tempdir will try."""

    dirlist = []

    # First, try the environment.
123
    for envname in 'TMPDIR', 'TEMP', 'TMP':
124 125 126 127 128
        dirname = _os.getenv(envname)
        if dirname: dirlist.append(dirname)

    # Failing that, try OS-specific locations.
    if _os.name == 'mac':
129
        try:
130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150
            refnum, dirid = _macfs.FindFolder(_MACFS.kOnSystemDisk,
                                              _MACFS.kTemporaryFolderType, 1)
            dirname = _macfs.FSSpec((refnum, dirid, '')).as_pathname()
            dirlist.append(dirname)
        except _macfs.error:
            pass
    elif _os.name == 'riscos':
        dirname = _os.getenv('Wimp$ScrapDir')
        if dirname: dirlist.append(dirname)
    elif _os.name == 'nt':
        dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
    else:
        dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])

    # As a last resort, the current directory.
    try:
        dirlist.append(_os.getcwd())
    except (AttributeError, _os.error):
        dirlist.append(_os.curdir)

    return dirlist
Tim Peters's avatar
Tim Peters committed
151

152 153
def _get_default_tempdir():
    """Calculate the default directory to use for temporary files.
154
    This routine should be called exactly once.
155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174

    We determine whether or not a candidate temp dir is usable by
    trying to create and write to a file in that directory.  If this
    is successful, the test file is deleted.  To prevent denial of
    service, the name of the test file must be randomized."""

    namer = _RandomNameSequence()
    dirlist = _candidate_tempdir_list()
    flags = _text_openflags

    for dir in dirlist:
        if dir != _os.curdir:
            dir = _os.path.normcase(_os.path.abspath(dir))
        # Try only a few names per directory.
        for seq in xrange(100):
            name = namer.next()
            filename = _os.path.join(dir, name)
            try:
                fd = _os.open(filename, flags, 0600)
                fp = _os.fdopen(fd, 'w')
Tim Peters's avatar
Tim Peters committed
175 176
                fp.write('blat')
                fp.close()
177 178 179 180 181 182 183 184 185
                _os.unlink(filename)
                del fp, fd
                return dir
            except (OSError, IOError), e:
                if e[0] != _errno.EEXIST:
                    break # no point trying more names in this directory
                pass
    raise IOError, (_errno.ENOENT,
                    ("No usable temporary directory found in %s" % dirlist))
186

187 188
_name_sequence = None

189 190
def _get_candidate_names():
    """Common setup sequence for all user-callable interfaces."""
191

192 193 194 195 196 197 198 199
    global _name_sequence
    if _name_sequence is None:
        _once_lock.acquire()
        try:
            if _name_sequence is None:
                _name_sequence = _RandomNameSequence()
        finally:
            _once_lock.release()
200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220
    return _name_sequence


def _mkstemp_inner(dir, pre, suf, flags):
    """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""

    names = _get_candidate_names()

    for seq in xrange(TMP_MAX):
        name = names.next()
        file = _os.path.join(dir, pre + name + suf)
        try:
            fd = _os.open(file, flags, 0600)
            _set_cloexec(fd)
            return (fd, file)
        except OSError, e:
            if e.errno == _errno.EEXIST:
                continue # try again
            raise

    raise IOError, (_errno.EEXIST, "No usable temporary file name found")
Tim Peters's avatar
Tim Peters committed
221

222 223

# User visible interfaces.
224

225
def gettempprefix():
226 227 228
    """Accessor for tempdir.template."""
    return template

229 230
tempdir = None

231 232
def gettempdir():
    """Accessor for tempdir.tempdir."""
233 234 235 236 237 238 239 240
    global tempdir
    if tempdir is None:
        _once_lock.acquire()
        try:
            if tempdir is None:
                tempdir = _get_default_tempdir()
        finally:
            _once_lock.release()
241 242
    return tempdir

243
def mkstemp(suffix="", prefix=template, dir=None, text=False):
244
    """mkstemp([suffix, [prefix, [dir, [text]]]])
245 246 247 248 249 250 251 252 253 254 255 256 257
    User-callable function to create and return a unique temporary
    file.  The return value is a pair (fd, name) where fd is the
    file descriptor returned by os.open, and name is the filename.

    If 'suffix' is specified, the file name will end with that suffix,
    otherwise there will be no suffix.

    If 'prefix' is specified, the file name will begin with that prefix,
    otherwise a default prefix is used.

    If 'dir' is specified, the file will be created in that directory,
    otherwise a default directory is used.

258 259 260
    If 'text' is specified and true, the file is opened in text
    mode.  Else (the default) the file is opened in binary mode.  On
    some operating systems, this makes no difference.
261 262 263 264 265

    The file is readable and writable only by the creating user ID.
    If the operating system uses permission bits to indicate whether a
    file is executable, the file is executable by no one. The file
    descriptor is not inherited by children of this process.
266

267
    Caller is responsible for deleting the file when done with it.
268 269
    """

270 271 272
    if dir is None:
        dir = gettempdir()

273
    if text:
274
        flags = _text_openflags
275 276
    else:
        flags = _bin_openflags
277

278
    return _mkstemp_inner(dir, prefix, suffix, flags)
Guido van Rossum's avatar
Guido van Rossum committed
279

280

281
def mkdtemp(suffix="", prefix=template, dir=None):
282 283 284 285
    """mkdtemp([suffix, [prefix, [dir]]])
    User-callable function to create and return a unique temporary
    directory.  The return value is the pathname of the directory.

286
    Arguments are as for mkstemp, except that the 'text' argument is
287 288 289 290 291 292 293 294
    not accepted.

    The directory is readable, writable, and searchable only by the
    creating user.

    Caller is responsible for deleting the directory when done with it.
    """

295 296 297
    if dir is None:
        dir = gettempdir()

298
    names = _get_candidate_names()
Tim Peters's avatar
Tim Peters committed
299

300 301 302 303 304
    for seq in xrange(TMP_MAX):
        name = names.next()
        file = _os.path.join(dir, prefix + name + suffix)
        try:
            _os.mkdir(file, 0700)
305
            return file
306 307 308 309
        except OSError, e:
            if e.errno == _errno.EEXIST:
                continue # try again
            raise
310

311
    raise IOError, (_errno.EEXIST, "No usable temporary directory name found")
312

313
def mktemp(suffix="", prefix=template, dir=None):
314 315 316
    """mktemp([suffix, [prefix, [dir]]])
    User-callable function to return a unique temporary file name.  The
    file is not created.
317

318
    Arguments are as for mkstemp, except that the 'text' argument is
319 320 321 322 323 324
    not accepted.

    This function is unsafe and should not be used.  The file name
    refers to a file that did not exist at some point, but by the time
    you get around to creating it, someone else may have beaten you to
    the punch.
325
    """
326

327 328 329
##    from warnings import warn as _warn
##    _warn("mktemp is a potential security risk to your program",
##          RuntimeWarning, stacklevel=2)
330

331 332 333
    if dir is None:
        dir = gettempdir()

334 335 336 337 338 339 340 341
    names = _get_candidate_names()
    for seq in xrange(TMP_MAX):
        name = names.next()
        file = _os.path.join(dir, prefix + name + suffix)
        if not _os.path.exists(file):
            return file

    raise IOError, (_errno.EEXIST, "No usable temporary filename found")
342

343 344 345 346 347 348 349
class _TemporaryFileWrapper:
    """Temporary file wrapper

    This class provides a wrapper around files opened for
    temporary use.  In particular, it seeks to automatically
    remove the file when it is no longer needed.
    """
350

351 352 353
    def __init__(self, file, name):
        self.file = file
        self.name = name
354
        self.close_called = False
355 356

    def __getattr__(self, name):
357 358
        file = self.__dict__['file']
        a = getattr(file, name)
359 360
        if type(a) != type(0):
            setattr(self, name, a)
361
        return a
362

363 364 365 366 367 368 369 370 371 372 373 374 375 376
    # NT provides delete-on-close as a primitive, so we don't need
    # the wrapper to do anything special.  We still use it so that
    # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
    if _os.name != 'nt':

        # Cache the unlinker so we don't get spurious errors at
        # shutdown when the module-level "os" is None'd out.  Note
        # that this must be referenced as self.unlink, because the
        # name TemporaryFileWrapper may also get None'd out before
        # __del__ is called.
        unlink = _os.unlink

        def close(self):
            if not self.close_called:
377
                self.close_called = True
378 379 380 381 382 383 384
                self.file.close()
                self.unlink(self.name)

        def __del__(self):
            self.close()

def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",
385
                       prefix=template, dir=None):
386 387 388 389 390 391 392 393 394 395 396
    """Create and return a temporary file.
    Arguments:
    'prefix', 'suffix', 'dir' -- as for mkstemp.
    'mode' -- the mode argument to os.fdopen (default "w+b").
    'bufsize' -- the buffer size argument to os.fdopen (default -1).
    The file is created as mkstemp() would do it.

    Returns a file object; the name of the file is accessible as
    file.name.  The file will be automatically deleted when it is
    closed.
    """
397

398 399 400
    if dir is None:
        dir = gettempdir()

401 402 403 404
    if 'b' in mode:
        flags = _bin_openflags
    else:
        flags = _text_openflags
405

406 407 408 409
    # Setting O_TEMPORARY in the flags causes the OS to delete
    # the file when it is closed.  This is only supported by Windows.
    if _os.name == 'nt':
        flags |= _os.O_TEMPORARY
410

411 412 413
    (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
    file = _os.fdopen(fd, mode, bufsize)
    return _TemporaryFileWrapper(file, name)
414

415 416 417
if _os.name != 'posix' or _os.sys.platform == 'cygwin':
    # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
    # while it is open.
418
    TemporaryFile = NamedTemporaryFile
419 420

else:
421
    def TemporaryFile(mode='w+b', bufsize=-1, suffix="",
422
                      prefix=template, dir=None):
423 424 425 426 427 428 429 430 431 432 433
        """Create and return a temporary file.
        Arguments:
        'prefix', 'suffix', 'directory' -- as for mkstemp.
        'mode' -- the mode argument to os.fdopen (default "w+b").
        'bufsize' -- the buffer size argument to os.fdopen (default -1).
        The file is created as mkstemp() would do it.

        Returns a file object.  The file has no name, and will cease to
        exist when it is closed.
        """

434 435 436
        if dir is None:
            dir = gettempdir()

437 438 439 440
        if 'b' in mode:
            flags = _bin_openflags
        else:
            flags = _text_openflags
441 442 443 444 445 446 447 448

        (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
        try:
            _os.unlink(name)
            return _os.fdopen(fd, mode, bufsize)
        except:
            _os.close(fd)
            raise