Lib/tempfile.py

   1 """Temporary files.
   2
   3 This module provides generic, low- and high-level interfaces for
   4 creating temporary files and directories.  The interfaces listed
   5 as "safe" just below can be used without fear of race conditions.
   6 Those listed as "unsafe" cannot, and are provided for backward
   7 compatibility only.
   8
   9 This 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     "mkstemp", "mkdtemp",                  # low level safe interfaces
  23     "mktemp",                              # deprecated unsafe interface
  24     "TMP_MAX", "gettempprefix",            # constants
  25     "tempdir", "gettempdir"
  26    ]
  27
  28
  29 # Imports.
  30
  31 import os as _os
  32 import errno as _errno
  33 from random import Random as _Random
  34
  35 if _os.name == 'mac':
  36     import macfs as _macfs
  37     import MACFS as _MACFS
  38
  39 try:
  40     import fcntl as _fcntl
  41     def _set_cloexec(fd):
  42         flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)
  43         if flags >= 0:
  44             # flags read successfully, modify
  45             flags |= _fcntl.FD_CLOEXEC
  46             _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)
  47 except (ImportError, AttributeError):
  48     def _set_cloexec(fd):
  49         pass
  50
  51 try:
  52     import thread as _thread
  53 except ImportError:
  54     import dummy_thread as _thread
  55 _allocate_lock = _thread.allocate_lock
  56
  57 _text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL
  58 if hasattr(_os, 'O_NOINHERIT'):
  59     _text_openflags |= _os.O_NOINHERIT
  60 if hasattr(_os, 'O_NOFOLLOW'):
  61     _text_openflags |= _os.O_NOFOLLOW
  62
  63 _bin_openflags = _text_openflags
  64 if hasattr(_os, 'O_BINARY'):
  65     _bin_openflags |= _os.O_BINARY
  66
  67 if hasattr(_os, 'TMP_MAX'):
  68     TMP_MAX = _os.TMP_MAX
  69 else:
  70     TMP_MAX = 10000
  71
  72 template = "tmp"
  73
  74 tempdir = None
  75
  76 # Internal routines.
  77
  78 _once_lock = _allocate_lock()
  79
  80 class _RandomNameSequence:
  81     """An instance of _RandomNameSequence generates an endless
  82     sequence of unpredictable strings which can safely be incorporated
  83     into file names.  Each string is six characters long.  Multiple
  84     threads can safely use the same instance at the same time.
  85
  86     _RandomNameSequence is an iterator."""
  87
  88     characters = ("abcdefghijklmnopqrstuvwxyz" +
  89                   "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
  90                   "0123456789-_")
  91
  92     def __init__(self):
  93         self.mutex = _allocate_lock()
  94         self.rng = _Random()
  95         self.normcase = _os.path.normcase
  96
  97     def __iter__(self):
  98         return self
  99
 100     def next(self):
 101         m = self.mutex
 102         c = self.characters
 103         choose = self.rng.choice
 104
 105         m.acquire()
 106         try:
 107             letters = [choose(c) for dummy in "123456"]
 108         finally:
 109             m.release()
 110
 111         return self.normcase(''.join(letters))
 112
 113 def _candidate_tempdir_list():
 114     """Generate a list of candidate temporary directories which
 115     _get_default_tempdir will try."""
 116
 117     dirlist = []
 118
 119     # First, try the environment.
 120     for envname in 'TMPDIR', 'TEMP', 'TMP':
 121         dirname = _os.getenv(envname)
 122         if dirname: dirlist.append(dirname)
 123
 124     # Failing that, try OS-specific locations.
 125     if _os.name == 'mac':
 126         try:
 127             refnum, dirid = _macfs.FindFolder(_MACFS.kOnSystemDisk,
 128                                               _MACFS.kTemporaryFolderType, 1)
 129             dirname = _macfs.FSSpec((refnum, dirid, '')).as_pathname()
 130             dirlist.append(dirname)
 131         except _macfs.error:
 132             pass
 133     elif _os.name == 'riscos':
 134         dirname = _os.getenv('Wimp$ScrapDir')
 135         if dirname: dirlist.append(dirname)
 136     elif _os.name == 'nt':
 137         dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
 138     else:
 139         dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])
 140
 141     # As a last resort, the current directory.
 142     try:
 143         dirlist.append(_os.getcwd())
 144     except (AttributeError, _os.error):
 145         dirlist.append(_os.curdir)
 146
 147     return dirlist
 148
 149 def _get_default_tempdir():
 150     """Calculate the default directory to use for temporary files.
 151     This routine should be called exactly once.
 152
 153     We determine whether or not a candidate temp dir is usable by
 154     trying to create and write to a file in that directory.  If this
 155     is successful, the test file is deleted.  To prevent denial of
 156     service, the name of the test file must be randomized."""
 157
 158     namer = _RandomNameSequence()
 159     dirlist = _candidate_tempdir_list()
 160     flags = _text_openflags
 161
 162     for dir in dirlist:
 163         if dir != _os.curdir:
 164             dir = _os.path.normcase(_os.path.abspath(dir))
 165         # Try only a few names per directory.
 166         for seq in xrange(100):
 167             name = namer.next()
 168             filename = _os.path.join(dir, name)
 169             try:
 170                 fd = _os.open(filename, flags, 0600)
 171                 fp = _os.fdopen(fd, 'w')
 172                 fp.write('blat')
 173                 fp.close()
 174                 _os.unlink(filename)
 175                 del fp, fd
 176                 return dir
 177             except (OSError, IOError), e:
 178                 if e[0] != _errno.EEXIST:
 179                     break # no point trying more names in this directory
 180                 pass
 181     raise IOError, (_errno.ENOENT,
 182                     ("No usable temporary directory found in %s" % dirlist))
 183
 184 _name_sequence = None
 185
 186 def _get_candidate_names():
 187     """Common setup sequence for all user-callable interfaces."""
 188
 189     global _name_sequence
 190     if _name_sequence is None:
 191         _once_lock.acquire()
 192         try:
 193             if _name_sequence is None:
 194                 _name_sequence = _RandomNameSequence()
 195         finally:
 196             _once_lock.release()
 197     return _name_sequence
 198
 199
 200 def _mkstemp_inner(dir, pre, suf, flags):
 201     """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
 202
 203     names = _get_candidate_names()
 204
 205     for seq in xrange(TMP_MAX):
 206         name = names.next()
 207         file = _os.path.join(dir, pre + name + suf)
 208         try:
 209             fd = _os.open(file, flags, 0600)
 210             _set_cloexec(fd)
 211             return (fd, file)
 212         except OSError, e:
 213             if e.errno == _errno.EEXIST:
 214                 continue # try again
 215             raise
 216
 217     raise IOError, (_errno.EEXIST, "No usable temporary file name found")
 218
 219
 220 # User visible interfaces.
 221
 222 def gettempprefix():
 223     """Accessor for tempdir.template."""
 224     return template
 225
 226 tempdir = None
 227
 228 def gettempdir():
 229     """Accessor for tempdir.tempdir."""
 230     global tempdir
 231     if tempdir is None:
 232         _once_lock.acquire()
 233         try:
 234             if tempdir is None:
 235                 tempdir = _get_default_tempdir()
 236         finally:
 237             _once_lock.release()
 238     return tempdir
 239
 240 def mkstemp(suffix="", prefix=template, dir=None, text=False):
 241     """mkstemp([suffix, [prefix, [dir, [text]]]])
 242     User-callable function to create and return a unique temporary
 243     file.  The return value is a pair (fd, name) where fd is the
 244     file descriptor returned by os.open, and name is the filename.
 245
 246     If 'suffix' is specified, the file name will end with that suffix,
 247     otherwise there will be no suffix.
 248
 249     If 'prefix' is specified, the file name will begin with that prefix,
 250     otherwise a default prefix is used.
 251
 252     If 'dir' is specified, the file will be created in that directory,
 253     otherwise a default directory is used.
 254
 255     If 'text' is specified and true, the file is opened in text
 256     mode.  Else (the default) the file is opened in binary mode.  On
 257     some operating systems, this makes no difference.
 258
 259     The file is readable and writable only by the creating user ID.
 260     If the operating system uses permission bits to indicate whether a
 261     file is executable, the file is executable by no one. The file
 262     descriptor is not inherited by children of this process.
 263
 264     Caller is responsible for deleting the file when done with it.
 265     """
 266
 267     if dir is None:
 268         dir = gettempdir()
 269
 270     if text:
 271         flags = _text_openflags
 272     else:
 273         flags = _bin_openflags
 274
 275     return _mkstemp_inner(dir, prefix, suffix, flags)
 276
 277
 278 def mkdtemp(suffix="", prefix=template, dir=None):
 279     """mkdtemp([suffix, [prefix, [dir]]])
 280     User-callable function to create and return a unique temporary
 281     directory.  The return value is the pathname of the directory.
 282
 283     Arguments are as for mkstemp, except that the 'text' argument is
 284     not accepted.
 285
 286     The directory is readable, writable, and searchable only by the
 287     creating user.
 288
 289     Caller is responsible for deleting the directory when done with it.
 290     """
 291
 292     if dir is None:
 293         dir = gettempdir()
 294
 295     names = _get_candidate_names()
 296
 297     for seq in xrange(TMP_MAX):
 298         name = names.next()
 299         file = _os.path.join(dir, prefix + name + suffix)
 300         try:
 301             _os.mkdir(file, 0700)
 302             return file
 303         except OSError, e:
 304             if e.errno == _errno.EEXIST:
 305                 continue # try again
 306             raise
 307
 308     raise IOError, (_errno.EEXIST, "No usable temporary directory name found")
 309
 310 def mktemp(suffix="", prefix=template, dir=None):
 311     """mktemp([suffix, [prefix, [dir]]])
 312     User-callable function to return a unique temporary file name.  The
 313     file is not created.
 314
 315     Arguments are as for mkstemp, except that the 'text' argument is
 316     not accepted.
 317
 318     This function is unsafe and should not be used.  The file name
 319     refers to a file that did not exist at some point, but by the time
 320     you get around to creating it, someone else may have beaten you to
 321     the punch.
 322     """
 323
 324 ##    from warnings import warn as _warn
 325 ##    _warn("mktemp is a potential security risk to your program",
 326 ##          RuntimeWarning, stacklevel=2)
 327
 328     if dir is None:
 329         dir = gettempdir()
 330
 331     names = _get_candidate_names()
 332     for seq in xrange(TMP_MAX):
 333         name = names.next()
 334         file = _os.path.join(dir, prefix + name + suffix)
 335         if not _os.path.exists(file):
 336             return file
 337
 338     raise IOError, (_errno.EEXIST, "No usable temporary filename found")
 339
 340 class _TemporaryFileWrapper:
 341     """Temporary file wrapper
 342
 343     This class provides a wrapper around files opened for
 344     temporary use.  In particular, it seeks to automatically
 345     remove the file when it is no longer needed.
 346     """
 347
 348     def __init__(self, file, name):
 349         self.file = file
 350         self.name = name
 351         self.close_called = False
 352
 353     def __getattr__(self, name):
 354         file = self.__dict__['file']
 355         a = getattr(file, name)
 356         if type(a) != type(0):
 357             setattr(self, name, a)
 358         return a
 359
 360     # NT provides delete-on-close as a primitive, so we don't need
 361     # the wrapper to do anything special.  We still use it so that
 362     # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
 363     if _os.name != 'nt':
 364
 365         # Cache the unlinker so we don't get spurious errors at
 366         # shutdown when the module-level "os" is None'd out.  Note
 367         # that this must be referenced as self.unlink, because the
 368         # name TemporaryFileWrapper may also get None'd out before
 369         # __del__ is called.
 370         unlink = _os.unlink
 371
 372         def close(self):
 373             if not self.close_called:
 374                 self.close_called = True
 375                 self.file.close()
 376                 self.unlink(self.name)
 377
 378         def __del__(self):
 379             self.close()
 380
 381 def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",
 382                        prefix=template, dir=None):
 383     """Create and return a temporary file.
 384     Arguments:
 385     'prefix', 'suffix', 'dir' -- as for mkstemp.
 386     'mode' -- the mode argument to os.fdopen (default "w+b").
 387     'bufsize' -- the buffer size argument to os.fdopen (default -1).
 388     The file is created as mkstemp() would do it.
 389
 390     Returns a file object; the name of the file is accessible as
 391     file.name.  The file will be automatically deleted when it is
 392     closed.
 393     """
 394
 395     if dir is None:
 396         dir = gettempdir()
 397
 398     if 'b' in mode:
 399         flags = _bin_openflags
 400     else:
 401         flags = _text_openflags
 402
 403     # Setting O_TEMPORARY in the flags causes the OS to delete
 404     # the file when it is closed.  This is only supported by Windows.
 405     if _os.name == 'nt':
 406         flags |= _os.O_TEMPORARY
 407
 408     (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
 409     file = _os.fdopen(fd, mode, bufsize)
 410     return _TemporaryFileWrapper(file, name)
 411
 412 if _os.name != 'posix' or _os.sys.platform == 'cygwin':
 413     # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
 414     # while it is open.
 415     TemporaryFile = NamedTemporaryFile
 416
 417 else:
 418     def TemporaryFile(mode='w+b', bufsize=-1, suffix="",
 419                       prefix=template, dir=None):
 420         """Create and return a temporary file.
 421         Arguments:
 422         'prefix', 'suffix', 'directory' -- as for mkstemp.
 423         'mode' -- the mode argument to os.fdopen (default "w+b").
 424         'bufsize' -- the buffer size argument to os.fdopen (default -1).
 425         The file is created as mkstemp() would do it.
 426
 427         Returns a file object.  The file has no name, and will cease to
 428         exist when it is closed.
 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         (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
 440         try:
 441             _os.unlink(name)
 442             return _os.fdopen(fd, mode, bufsize)
 443         except:
 444             _os.close(fd)
 445             raise