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 Carbon.Folder as _Folder
  37     import Carbon.Folders as _Folders
  38
  39 try:
  40     import fcntl as _fcntl
  41 except ImportError:
  42     def _set_cloexec(fd):
  43         pass
  44 else:
  45     def _set_cloexec(fd):
  46         try:
  47             flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)
  48         except IOError:
  49             pass
  50         else:
  51             # flags read successfully, modify
  52             flags |= _fcntl.FD_CLOEXEC
  53             _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)
  54
  55
  56 try:
  57     import thread as _thread
  58 except ImportError:
  59     import dummy_thread as _thread
  60 _allocate_lock = _thread.allocate_lock
  61
  62 _text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL
  63 if hasattr(_os, 'O_NOINHERIT'):
  64     _text_openflags |= _os.O_NOINHERIT
  65 if hasattr(_os, 'O_NOFOLLOW'):
  66     _text_openflags |= _os.O_NOFOLLOW
  67
  68 _bin_openflags = _text_openflags
  69 if hasattr(_os, 'O_BINARY'):
  70     _bin_openflags |= _os.O_BINARY
  71
  72 if hasattr(_os, 'TMP_MAX'):
  73     TMP_MAX = _os.TMP_MAX
  74 else:
  75     TMP_MAX = 10000
  76
  77 template = "tmp"
  78
  79 tempdir = None
  80
  81 # Internal routines.
  82
  83 _once_lock = _allocate_lock()
  84
  85 if hasattr(_os, "lstat"):
  86     _stat = _os.lstat
  87 elif hasattr(_os, "stat"):
  88     _stat = _os.stat
  89 else:
  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
  99 def _exists(fn):
 100     try:
 101         _stat(fn)
 102     except _os.error:
 103         return False
 104     else:
 105         return True
 106
 107 class _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
 140 def _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 == 'mac':
 153         try:
 154             fsr = _Folder.FSFindFolder(_Folders.kOnSystemDisk,
 155                                               _Folders.kTemporaryFolderType, 1)
 156             dirname = fsr.as_pathname()
 157             dirlist.append(dirname)
 158         except _Folder.error:
 159             pass
 160     elif _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
 176 def _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, 0600)
 198                 fp = _os.fdopen(fd, 'w')
 199                 fp.write('blat')
 200                 fp.close()
 201                 _os.unlink(filename)
 202                 del fp, fd
 203                 return dir
 204             except (OSError, IOError), e:
 205                 if e[0] != _errno.EEXIST:
 206                     break # no point trying more names in this directory
 207                 pass
 208     raise IOError, (_errno.ENOENT,
 209                     ("No usable temporary directory found in %s" % dirlist))
 210
 211 _name_sequence = None
 212
 213 def _get_candidate_names():
 214     """Common setup sequence for all user-callable interfaces."""
 215
 216     global _name_sequence
 217     if _name_sequence is None:
 218         _once_lock.acquire()
 219         try:
 220             if _name_sequence is None:
 221                 _name_sequence = _RandomNameSequence()
 222         finally:
 223             _once_lock.release()
 224     return _name_sequence
 225
 226
 227 def _mkstemp_inner(dir, pre, suf, flags):
 228     """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
 229
 230     names = _get_candidate_names()
 231
 232     for seq in xrange(TMP_MAX):
 233         name = names.next()
 234         file = _os.path.join(dir, pre + name + suf)
 235         try:
 236             fd = _os.open(file, flags, 0600)
 237             _set_cloexec(fd)
 238             return (fd, _os.path.abspath(file))
 239         except OSError, e:
 240             if e.errno == _errno.EEXIST:
 241                 continue # try again
 242             raise
 243
 244     raise IOError, (_errno.EEXIST, "No usable temporary file name found")
 245
 246
 247 # User visible interfaces.
 248
 249 def gettempprefix():
 250     """Accessor for tempdir.template."""
 251     return template
 252
 253 tempdir = None
 254
 255 def gettempdir():
 256     """Accessor for tempdir.tempdir."""
 257     global tempdir
 258     if tempdir is None:
 259         _once_lock.acquire()
 260         try:
 261             if tempdir is None:
 262                 tempdir = _get_default_tempdir()
 263         finally:
 264             _once_lock.release()
 265     return tempdir
 266
 267 def mkstemp(suffix="", prefix=template, dir=None, text=False):
 268     """mkstemp([suffix, [prefix, [dir, [text]]]])
 269     User-callable function to create and return a unique temporary
 270     file.  The return value is a pair (fd, name) where fd is the
 271     file descriptor returned by os.open, and name is the filename.
 272
 273     If 'suffix' is specified, the file name will end with that suffix,
 274     otherwise there will be no suffix.
 275
 276     If 'prefix' is specified, the file name will begin with that prefix,
 277     otherwise a default prefix is used.
 278
 279     If 'dir' is specified, the file will be created in that directory,
 280     otherwise a default directory is used.
 281
 282     If 'text' is specified and true, the file is opened in text
 283     mode.  Else (the default) the file is opened in binary mode.  On
 284     some operating systems, this makes no difference.
 285
 286     The file is readable and writable only by the creating user ID.
 287     If the operating system uses permission bits to indicate whether a
 288     file is executable, the file is executable by no one. The file
 289     descriptor is not inherited by children of this process.
 290
 291     Caller is responsible for deleting the file when done with it.
 292     """
 293
 294     if dir is None:
 295         dir = gettempdir()
 296
 297     if text:
 298         flags = _text_openflags
 299     else:
 300         flags = _bin_openflags
 301
 302     return _mkstemp_inner(dir, prefix, suffix, flags)
 303
 304
 305 def mkdtemp(suffix="", prefix=template, dir=None):
 306     """mkdtemp([suffix, [prefix, [dir]]])
 307     User-callable function to create and return a unique temporary
 308     directory.  The return value is the pathname of the directory.
 309
 310     Arguments are as for mkstemp, except that the 'text' argument is
 311     not accepted.
 312
 313     The directory is readable, writable, and searchable only by the
 314     creating user.
 315
 316     Caller is responsible for deleting the directory when done with it.
 317     """
 318
 319     if dir is None:
 320         dir = gettempdir()
 321
 322     names = _get_candidate_names()
 323
 324     for seq in xrange(TMP_MAX):
 325         name = names.next()
 326         file = _os.path.join(dir, prefix + name + suffix)
 327         try:
 328             _os.mkdir(file, 0700)
 329             return file
 330         except OSError, e:
 331             if e.errno == _errno.EEXIST:
 332                 continue # try again
 333             raise
 334
 335     raise IOError, (_errno.EEXIST, "No usable temporary directory name found")
 336
 337 def mktemp(suffix="", prefix=template, dir=None):
 338     """mktemp([suffix, [prefix, [dir]]])
 339     User-callable function to return a unique temporary file name.  The
 340     file is not created.
 341
 342     Arguments are as for mkstemp, except that the 'text' argument is
 343     not accepted.
 344
 345     This function is unsafe and should not be used.  The file name
 346     refers to a file that did not exist at some point, but by the time
 347     you get around to creating it, someone else may have beaten you to
 348     the punch.
 349     """
 350
 351 ##    from warnings import warn as _warn
 352 ##    _warn("mktemp is a potential security risk to your program",
 353 ##          RuntimeWarning, stacklevel=2)
 354
 355     if dir is None:
 356         dir = gettempdir()
 357
 358     names = _get_candidate_names()
 359     for seq in xrange(TMP_MAX):
 360         name = names.next()
 361         file = _os.path.join(dir, prefix + name + suffix)
 362         if not _exists(file):
 363             return file
 364
 365     raise IOError, (_errno.EEXIST, "No usable temporary filename found")
 366
 367 class _TemporaryFileWrapper:
 368     """Temporary file wrapper
 369
 370     This class provides a wrapper around files opened for
 371     temporary use.  In particular, it seeks to automatically
 372     remove the file when it is no longer needed.
 373     """
 374
 375     def __init__(self, file, name):
 376         self.file = file
 377         self.name = name
 378         self.close_called = False
 379
 380     def __getattr__(self, name):
 381         file = self.__dict__['file']
 382         a = getattr(file, name)
 383         if type(a) != type(0):
 384             setattr(self, name, a)
 385         return a
 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
 392         # Cache the unlinker so we don't get spurious errors at
 393         # shutdown when the module-level "os" is None'd out.  Note
 394         # that this must be referenced as self.unlink, because the
 395         # name TemporaryFileWrapper may also get None'd out before
 396         # __del__ is called.
 397         unlink = _os.unlink
 398
 399         def close(self):
 400             if not self.close_called:
 401                 self.close_called = True
 402                 self.file.close()
 403                 self.unlink(self.name)
 404
 405         def __del__(self):
 406             self.close()
 407
 408 def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",
 409                        prefix=template, dir=None):
 410     """Create and return a temporary file.
 411     Arguments:
 412     'prefix', 'suffix', 'dir' -- as for mkstemp.
 413     'mode' -- the mode argument to os.fdopen (default "w+b").
 414     'bufsize' -- the buffer size argument to os.fdopen (default -1).
 415     The file is created as mkstemp() would do it.
 416
 417     Returns an object with a file-like interface; the name of the file
 418     is accessible as file.name.  The file will be automatically deleted
 419     when it is closed.
 420     """
 421
 422     if dir is None:
 423         dir = gettempdir()
 424
 425     if 'b' in mode:
 426         flags = _bin_openflags
 427     else:
 428         flags = _text_openflags
 429
 430     # Setting O_TEMPORARY in the flags causes the OS to delete
 431     # the file when it is closed.  This is only supported by Windows.
 432     if _os.name == 'nt':
 433         flags |= _os.O_TEMPORARY
 434
 435     (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
 436     file = _os.fdopen(fd, mode, bufsize)
 437     return _TemporaryFileWrapper(file, name)
 438
 439 if _os.name != 'posix' or _os.sys.platform == 'cygwin':
 440     # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
 441     # while it is open.
 442     TemporaryFile = NamedTemporaryFile
 443
 444 else:
 445     def TemporaryFile(mode='w+b', bufsize=-1, suffix="",
 446                       prefix=template, dir=None):
 447         """Create and return a temporary file.
 448         Arguments:
 449         'prefix', 'suffix', 'directory' -- as for mkstemp.
 450         'mode' -- the mode argument to os.fdopen (default "w+b").
 451         'bufsize' -- the buffer size argument to os.fdopen (default -1).
 452         The file is created as mkstemp() would do it.
 453
 454         Returns an object with a file-like interface.  The file has no
 455         name, and will cease to exist when it is closed.
 456         """
 457
 458         if dir is None:
 459             dir = gettempdir()
 460
 461         if 'b' in mode:
 462             flags = _bin_openflags
 463         else:
 464             flags = _text_openflags
 465
 466         (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
 467         try:
 468             _os.unlink(name)
 469             return _os.fdopen(fd, mode, bufsize)
 470         except:
 471             _os.close(fd)
 472             raise