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 if _os.name == 'mac':
  37     import Carbon.Folder as _Folder
  38     import Carbon.Folders as _Folders
  39
  40 try:
  41     from cStringIO import StringIO as _StringIO
  42 except ImportError:
  43     from StringIO import StringIO as _StringIO
  44
  45 try:
  46     import fcntl as _fcntl
  47 except ImportError:
  48     def _set_cloexec(fd):
  49         pass
  50 else:
  51     def _set_cloexec(fd):
  52         try:
  53             flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)
  54         except IOError:
  55             pass
  56         else:
  57             # flags read successfully, modify
  58             flags |= _fcntl.FD_CLOEXEC
  59             _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)
  60
  61
  62 try:
  63     import thread as _thread
  64 except ImportError:
  65     import dummy_thread as _thread
  66 _allocate_lock = _thread.allocate_lock
  67
  68 _text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL
  69 if hasattr(_os, 'O_NOINHERIT'):
  70     _text_openflags |= _os.O_NOINHERIT
  71 if hasattr(_os, 'O_NOFOLLOW'):
  72     _text_openflags |= _os.O_NOFOLLOW
  73
  74 _bin_openflags = _text_openflags
  75 if hasattr(_os, 'O_BINARY'):
  76     _bin_openflags |= _os.O_BINARY
  77
  78 if hasattr(_os, 'TMP_MAX'):
  79     TMP_MAX = _os.TMP_MAX
  80 else:
  81     TMP_MAX = 10000
  82
  83 template = "tmp"
  84
  85 # Internal routines.
  86
  87 _once_lock = _allocate_lock()
  88
  89 if hasattr(_os, "lstat"):
  90     _stat = _os.lstat
  91 elif hasattr(_os, "stat"):
  92     _stat = _os.stat
  93 else:
  94     # Fallback.  All we need is something that raises os.error if the
  95     # file doesn't exist.
  96     def _stat(fn):
  97         try:
  98             f = open(fn)
  99         except IOError:
 100             raise _os.error
 101         f.close()
 102
 103 def _exists(fn):
 104     try:
 105         _stat(fn)
 106     except _os.error:
 107         return False
 108     else:
 109         return True
 110
 111 class _RandomNameSequence:
 112     """An instance of _RandomNameSequence generates an endless
 113     sequence of unpredictable strings which can safely be incorporated
 114     into file names.  Each string is six characters long.  Multiple
 115     threads can safely use the same instance at the same time.
 116
 117     _RandomNameSequence is an iterator."""
 118
 119     characters = ("abcdefghijklmnopqrstuvwxyz" +
 120                   "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
 121                   "0123456789_")
 122
 123     def __init__(self):
 124         self.mutex = _allocate_lock()
 125         self.rng = _Random()
 126         self.normcase = _os.path.normcase
 127
 128     def __iter__(self):
 129         return self
 130
 131     def next(self):
 132         m = self.mutex
 133         c = self.characters
 134         choose = self.rng.choice
 135
 136         m.acquire()
 137         try:
 138             letters = [choose(c) for dummy in "123456"]
 139         finally:
 140             m.release()
 141
 142         return self.normcase(''.join(letters))
 143
 144 def _candidate_tempdir_list():
 145     """Generate a list of candidate temporary directories which
 146     _get_default_tempdir will try."""
 147
 148     dirlist = []
 149
 150     # First, try the environment.
 151     for envname in 'TMPDIR', 'TEMP', 'TMP':
 152         dirname = _os.getenv(envname)
 153         if dirname: dirlist.append(dirname)
 154
 155     # Failing that, try OS-specific locations.
 156     if _os.name == 'mac':
 157         try:
 158             fsr = _Folder.FSFindFolder(_Folders.kOnSystemDisk,
 159                                               _Folders.kTemporaryFolderType, 1)
 160             dirname = fsr.as_pathname()
 161             dirlist.append(dirname)
 162         except _Folder.error:
 163             pass
 164     elif _os.name == 'riscos':
 165         dirname = _os.getenv('Wimp$ScrapDir')
 166         if dirname: dirlist.append(dirname)
 167     elif _os.name == 'nt':
 168         dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
 169     else:
 170         dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])
 171
 172     # As a last resort, the current directory.
 173     try:
 174         dirlist.append(_os.getcwd())
 175     except (AttributeError, _os.error):
 176         dirlist.append(_os.curdir)
 177
 178     return dirlist
 179
 180 def _get_default_tempdir():
 181     """Calculate the default directory to use for temporary files.
 182     This routine should be called exactly once.
 183
 184     We determine whether or not a candidate temp dir is usable by
 185     trying to create and write to a file in that directory.  If this
 186     is successful, the test file is deleted.  To prevent denial of
 187     service, the name of the test file must be randomized."""
 188
 189     namer = _RandomNameSequence()
 190     dirlist = _candidate_tempdir_list()
 191     flags = _text_openflags
 192
 193     for dir in dirlist:
 194         if dir != _os.curdir:
 195             dir = _os.path.normcase(_os.path.abspath(dir))
 196         # Try only a few names per directory.
 197         for seq in xrange(100):
 198             name = namer.next()
 199             filename = _os.path.join(dir, name)
 200             try:
 201                 fd = _os.open(filename, flags, 0600)
 202                 fp = _os.fdopen(fd, 'w')
 203                 fp.write('blat')
 204                 fp.close()
 205                 _os.unlink(filename)
 206                 del fp, fd
 207                 return dir
 208             except (OSError, IOError), e:
 209                 if e[0] != _errno.EEXIST:
 210                     break # no point trying more names in this directory
 211                 pass
 212     raise IOError, (_errno.ENOENT,
 213                     ("No usable temporary directory found in %s" % dirlist))
 214
 215 _name_sequence = None
 216
 217 def _get_candidate_names():
 218     """Common setup sequence for all user-callable interfaces."""
 219
 220     global _name_sequence
 221     if _name_sequence is None:
 222         _once_lock.acquire()
 223         try:
 224             if _name_sequence is None:
 225                 _name_sequence = _RandomNameSequence()
 226         finally:
 227             _once_lock.release()
 228     return _name_sequence
 229
 230
 231 def _mkstemp_inner(dir, pre, suf, flags):
 232     """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
 233
 234     names = _get_candidate_names()
 235
 236     for seq in xrange(TMP_MAX):
 237         name = names.next()
 238         file = _os.path.join(dir, pre + name + suf)
 239         try:
 240             fd = _os.open(file, flags, 0600)
 241             _set_cloexec(fd)
 242             return (fd, _os.path.abspath(file))
 243         except OSError, e:
 244             if e.errno == _errno.EEXIST:
 245                 continue # try again
 246             raise
 247
 248     raise IOError, (_errno.EEXIST, "No usable temporary file name found")
 249
 250
 251 # User visible interfaces.
 252
 253 def gettempprefix():
 254     """Accessor for tempdir.template."""
 255     return template
 256
 257 tempdir = None
 258
 259 def gettempdir():
 260     """Accessor for tempfile.tempdir."""
 261     global tempdir
 262     if tempdir is None:
 263         _once_lock.acquire()
 264         try:
 265             if tempdir is None:
 266                 tempdir = _get_default_tempdir()
 267         finally:
 268             _once_lock.release()
 269     return tempdir
 270
 271 def mkstemp(suffix="", prefix=template, dir=None, text=False):
 272     """User-callable function to create and return a unique temporary
 273     file.  The return value is a pair (fd, name) where fd is the
 274     file descriptor returned by os.open, and name is the filename.
 275
 276     If 'suffix' is specified, the file name will end with that suffix,
 277     otherwise there will be no suffix.
 278
 279     If 'prefix' is specified, the file name will begin with that prefix,
 280     otherwise a default prefix is used.
 281
 282     If 'dir' is specified, the file will be created in that directory,
 283     otherwise a default directory is used.
 284
 285     If 'text' is specified and true, the file is opened in text
 286     mode.  Else (the default) the file is opened in binary mode.  On
 287     some operating systems, this makes no difference.
 288
 289     The file is readable and writable only by the creating user ID.
 290     If the operating system uses permission bits to indicate whether a
 291     file is executable, the file is executable by no one. The file
 292     descriptor is not inherited by children of this process.
 293
 294     Caller is responsible for deleting the file when done with it.
 295     """
 296
 297     if dir is None:
 298         dir = gettempdir()
 299
 300     if text:
 301         flags = _text_openflags
 302     else:
 303         flags = _bin_openflags
 304
 305     return _mkstemp_inner(dir, prefix, suffix, flags)
 306
 307
 308 def mkdtemp(suffix="", prefix=template, dir=None):
 309     """User-callable function to create and return a unique temporary
 310     directory.  The return value is the pathname of the directory.
 311
 312     Arguments are as for mkstemp, except that the 'text' argument is
 313     not accepted.
 314
 315     The directory is readable, writable, and searchable only by the
 316     creating user.
 317
 318     Caller is responsible for deleting the directory when done with it.
 319     """
 320
 321     if dir is None:
 322         dir = gettempdir()
 323
 324     names = _get_candidate_names()
 325
 326     for seq in xrange(TMP_MAX):
 327         name = names.next()
 328         file = _os.path.join(dir, prefix + name + suffix)
 329         try:
 330             _os.mkdir(file, 0700)
 331             return file
 332         except OSError, e:
 333             if e.errno == _errno.EEXIST:
 334                 continue # try again
 335             raise
 336
 337     raise IOError, (_errno.EEXIST, "No usable temporary directory name found")
 338
 339 def mktemp(suffix="", prefix=template, dir=None):
 340     """User-callable function to return a unique temporary file name.  The
 341     file is not created.
 342
 343     Arguments are as for mkstemp, except that the 'text' argument is
 344     not accepted.
 345
 346     This function is unsafe and should not be used.  The file name
 347     refers to a file that did not exist at some point, but by the time
 348     you get around to creating it, someone else may have beaten you to
 349     the punch.
 350     """
 351
 352 ##    from warnings import warn as _warn
 353 ##    _warn("mktemp is a potential security risk to your program",
 354 ##          RuntimeWarning, stacklevel=2)
 355
 356     if dir is None:
 357         dir = gettempdir()
 358
 359     names = _get_candidate_names()
 360     for seq in xrange(TMP_MAX):
 361         name = names.next()
 362         file = _os.path.join(dir, prefix + name + suffix)
 363         if not _exists(file):
 364             return file
 365
 366     raise IOError, (_errno.EEXIST, "No usable temporary filename found")
 367
 368
 369 class _TemporaryFileWrapper:
 370     """Temporary file wrapper
 371
 372     This class provides a wrapper around files opened for
 373     temporary use.  In particular, it seeks to automatically
 374     remove the file when it is no longer needed.
 375     """
 376
 377     def __init__(self, file, name, delete=True):
 378         self.file = file
 379         self.name = name
 380         self.close_called = False
 381         self.delete = delete
 382
 383     def __getattr__(self, name):
 384         # Attribute lookups are delegated to the underlying file
 385         # and cached for non-numeric results
 386         # (i.e. methods are cached, closed and friends are not)
 387         file = self.__dict__['file']
 388         a = getattr(file, name)
 389         if not issubclass(type(a), type(0)):
 390             setattr(self, name, a)
 391         return a
 392
 393     # The underlying __enter__ method returns the wrong object
 394     # (self.file) so override it to return the wrapper
 395     def __enter__(self):
 396         self.file.__enter__()
 397         return self
 398
 399     # NT provides delete-on-close as a primitive, so we don't need
 400     # the wrapper to do anything special.  We still use it so that
 401     # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
 402     if _os.name != 'nt':
 403         # Cache the unlinker so we don't get spurious errors at
 404         # shutdown when the module-level "os" is None'd out.  Note
 405         # that this must be referenced as self.unlink, because the
 406         # name TemporaryFileWrapper may also get None'd out before
 407         # __del__ is called.
 408         unlink = _os.unlink
 409
 410         def close(self):
 411             if not self.close_called:
 412                 self.close_called = True
 413                 self.file.close()
 414                 if self.delete:
 415                     self.unlink(self.name)
 416
 417         def __del__(self):
 418             self.close()
 419
 420         # Need to trap __exit__ as well to ensure the file gets
 421         # deleted when used in a with statement
 422         def __exit__(self, exc, value, tb):
 423             result = self.file.__exit__(exc, value, tb)
 424             self.close()
 425             return result
 426
 427
 428 def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",
 429                        prefix=template, dir=None, delete=True):
 430     """Create and return a temporary file.
 431     Arguments:
 432     'prefix', 'suffix', 'dir' -- as for mkstemp.
 433     'mode' -- the mode argument to os.fdopen (default "w+b").
 434     'bufsize' -- the buffer size argument to os.fdopen (default -1).
 435     'delete' -- whether the file is deleted on close (default True).
 436     The file is created as mkstemp() would do it.
 437
 438     Returns an object with a file-like interface; the name of the file
 439     is accessible as file.name.  The file will be automatically deleted
 440     when it is closed unless the 'delete' argument is set to False.
 441     """
 442
 443     if dir is None:
 444         dir = gettempdir()
 445
 446     if 'b' in mode:
 447         flags = _bin_openflags
 448     else:
 449         flags = _text_openflags
 450
 451     # Setting O_TEMPORARY in the flags causes the OS to delete
 452     # the file when it is closed.  This is only supported by Windows.
 453     if _os.name == 'nt' and delete:
 454         flags |= _os.O_TEMPORARY
 455
 456     (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
 457     file = _os.fdopen(fd, mode, bufsize)
 458     return _TemporaryFileWrapper(file, name, delete)
 459
 460 if _os.name != 'posix' or _os.sys.platform == 'cygwin':
 461     # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
 462     # while it is open.
 463     TemporaryFile = NamedTemporaryFile
 464
 465 else:
 466     def TemporaryFile(mode='w+b', bufsize=-1, suffix="",
 467                       prefix=template, dir=None):
 468         """Create and return a temporary file.
 469         Arguments:
 470         'prefix', 'suffix', 'dir' -- as for mkstemp.
 471         'mode' -- the mode argument to os.fdopen (default "w+b").
 472         'bufsize' -- the buffer size argument to os.fdopen (default -1).
 473         The file is created as mkstemp() would do it.
 474
 475         Returns an object with a file-like interface.  The file has no
 476         name, and will cease to exist when it is closed.
 477         """
 478
 479         if dir is None:
 480             dir = gettempdir()
 481
 482         if 'b' in mode:
 483             flags = _bin_openflags
 484         else:
 485             flags = _text_openflags
 486
 487         (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
 488         try:
 489             _os.unlink(name)
 490             return _os.fdopen(fd, mode, bufsize)
 491         except:
 492             _os.close(fd)
 493             raise
 494
 495 class SpooledTemporaryFile:
 496     """Temporary file wrapper, specialized to switch from
 497     StringIO to a real file when it exceeds a certain size or
 498     when a fileno is needed.
 499     """
 500     _rolled = False
 501
 502     def __init__(self, max_size=0, mode='w+b', bufsize=-1,
 503                  suffix="", prefix=template, dir=None):
 504         self._file = _StringIO()
 505         self._max_size = max_size
 506         self._rolled = False
 507         self._TemporaryFileArgs = (mode, bufsize, suffix, prefix, dir)
 508
 509     def _check(self, file):
 510         if self._rolled: return
 511         max_size = self._max_size
 512         if max_size and file.tell() > max_size:
 513             self.rollover()
 514
 515     def rollover(self):
 516         if self._rolled: return
 517         file = self._file
 518         newfile = self._file = TemporaryFile(*self._TemporaryFileArgs)
 519         del self._TemporaryFileArgs
 520
 521         newfile.write(file.getvalue())
 522         newfile.seek(file.tell(), 0)
 523
 524         self._rolled = True
 525
 526     # The method caching trick from NamedTemporaryFile
 527     # won't work here, because _file may change from a
 528     # _StringIO instance to a real file. So we list
 529     # all the methods directly.
 530
 531     # Context management protocol
 532     def __enter__(self):
 533         if self._file.closed:
 534             raise ValueError("Cannot enter context with closed file")
 535         return self
 536
 537     def __exit__(self, exc, value, tb):
 538         self._file.close()
 539
 540     # file protocol
 541     def __iter__(self):
 542         return self._file.__iter__()
 543
 544     def close(self):
 545         self._file.close()
 546
 547     @property
 548     def closed(self):
 549         return self._file.closed
 550
 551     @property
 552     def encoding(self):
 553         return self._file.encoding
 554
 555     def fileno(self):
 556         self.rollover()
 557         return self._file.fileno()
 558
 559     def flush(self):
 560         self._file.flush()
 561
 562     def isatty(self):
 563         return self._file.isatty()
 564
 565     @property
 566     def mode(self):
 567         return self._file.mode
 568
 569     @property
 570     def name(self):
 571         return self._file.name
 572
 573     @property
 574     def newlines(self):
 575         return self._file.newlines
 576
 577     def next(self):
 578         return self._file.next
 579
 580     def read(self, *args):
 581         return self._file.read(*args)
 582
 583     def readline(self, *args):
 584         return self._file.readline(*args)
 585
 586     def readlines(self, *args):
 587         return self._file.readlines(*args)
 588
 589     def seek(self, *args):
 590         self._file.seek(*args)
 591
 592     @property
 593     def softspace(self):
 594         return self._file.softspace
 595
 596     def tell(self):
 597         return self._file.tell()
 598
 599     def truncate(self):
 600         self._file.truncate()
 601
 602     def write(self, s):
 603         file = self._file
 604         rv = file.write(s)
 605         self._check(file)
 606         return rv
 607
 608     def writelines(self, iterable):
 609         file = self._file
 610         rv = file.writelines(iterable)
 611         self._check(file)
 612         return rv
 613
 614     def xreadlines(self, *args):
 615         return self._file.xreadlines(*args)