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     "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
  32 import os as _os
  33 import errno as _errno
  34 from random import Random as _Random
  35
  36 try:
  37     from cStringIO import StringIO as _StringIO
  38 except ImportError:
  39     from StringIO import StringIO as _StringIO
  40
  41 try:
  42     import fcntl as _fcntl
  43 except ImportError:
  44     def _set_cloexec(fd):
  45         pass
  46 else:
  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
  58 try:
  59     import thread as _thread
  60 except 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
  65 if hasattr(_os, 'O_NOINHERIT'):
  66     _text_openflags |= _os.O_NOINHERIT
  67 if hasattr(_os, 'O_NOFOLLOW'):
  68     _text_openflags |= _os.O_NOFOLLOW
  69
  70 _bin_openflags = _text_openflags
  71 if hasattr(_os, 'O_BINARY'):
  72     _bin_openflags |= _os.O_BINARY
  73
  74 if hasattr(_os, 'TMP_MAX'):
  75     TMP_MAX = _os.TMP_MAX
  76 else:
  77     TMP_MAX = 10000
  78
  79 template = "tmp"
  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 == '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
 168 def _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
 205 def _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
 219 def _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
 241 def gettempprefix():
 242     """Accessor for tempdir.template."""
 243     return template
 244
 245 tempdir = None
 246
 247 def 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
 259 def 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
 296 def 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
 327 def 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
 357 class _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
 416 def 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
 448 if _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
 453 else:
 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
 483 class 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)