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