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:
  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 tempdir = None
  86
  87 # Internal routines.
  88
  89 _once_lock = _allocate_lock()
  90
  91 if hasattr(_os, "lstat"):
  92     _stat = _os.lstat
  93 elif hasattr(_os, "stat"):
  94     _stat = _os.stat
  95 else:
  96     # Fallback.  All we need is something that raises os.error if the
  97     # file doesn't exist.
  98     def _stat(fn):
  99         try:
 100             f = open(fn)
 101         except IOError:
 102             raise _os.error
 103         f.close()
 104
 105 def _exists(fn):
 106     try:
 107         _stat(fn)
 108     except _os.error:
 109         return False
 110     else:
 111         return True
 112
 113 class _RandomNameSequence:
 114     """An instance of _RandomNameSequence generates an endless
 115     sequence of unpredictable strings which can safely be incorporated
 116     into file names.  Each string is six characters long.  Multiple
 117     threads can safely use the same instance at the same time.
 118
 119     _RandomNameSequence is an iterator."""
 120
 121     characters = ("abcdefghijklmnopqrstuvwxyz" +
 122                   "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
 123                   "0123456789_")
 124
 125     def __init__(self):
 126         self.mutex = _allocate_lock()
 127         self.rng = _Random()
 128         self.normcase = _os.path.normcase
 129
 130     def __iter__(self):
 131         return self
 132
 133     def next(self):
 134         m = self.mutex
 135         c = self.characters
 136         choose = self.rng.choice
 137
 138         m.acquire()
 139         try:
 140             letters = [choose(c) for dummy in "123456"]
 141         finally:
 142             m.release()
 143
 144         return self.normcase(''.join(letters))
 145
 146 def _candidate_tempdir_list():
 147     """Generate a list of candidate temporary directories which
 148     _get_default_tempdir will try."""
 149
 150     dirlist = []
 151
 152     # First, try the environment.
 153     for envname in 'TMPDIR', 'TEMP', 'TMP':
 154         dirname = _os.getenv(envname)
 155         if dirname: dirlist.append(dirname)
 156
 157     # Failing that, try OS-specific locations.
 158     if _os.name == 'mac':
 159         try:
 160             fsr = _Folder.FSFindFolder(_Folders.kOnSystemDisk,
 161                                               _Folders.kTemporaryFolderType, 1)
 162             dirname = fsr.as_pathname()
 163             dirlist.append(dirname)
 164         except _Folder.error:
 165             pass
 166     elif _os.name == 'riscos':
 167         dirname = _os.getenv('Wimp$ScrapDir')
 168         if dirname: dirlist.append(dirname)
 169     elif _os.name == 'nt':
 170         dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
 171     else:
 172         dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])
 173
 174     # As a last resort, the current directory.
 175     try:
 176         dirlist.append(_os.getcwd())
 177     except (AttributeError, _os.error):
 178         dirlist.append(_os.curdir)
 179
 180     return dirlist
 181
 182 def _get_default_tempdir():
 183     """Calculate the default directory to use for temporary files.
 184     This routine should be called exactly once.
 185
 186     We determine whether or not a candidate temp dir is usable by
 187     trying to create and write to a file in that directory.  If this
 188     is successful, the test file is deleted.  To prevent denial of
 189     service, the name of the test file must be randomized."""
 190
 191     namer = _RandomNameSequence()
 192     dirlist = _candidate_tempdir_list()
 193     flags = _text_openflags
 194
 195     for dir in dirlist:
 196         if dir != _os.curdir:
 197             dir = _os.path.normcase(_os.path.abspath(dir))
 198         # Try only a few names per directory.
 199         for seq in xrange(100):
 200             name = namer.next()
 201             filename = _os.path.join(dir, name)
 202             try:
 203                 fd = _os.open(filename, flags, 0600)
 204                 fp = _os.fdopen(fd, 'w')
 205                 fp.write('blat')
 206                 fp.close()
 207                 _os.unlink(filename)
 208                 del fp, fd
 209                 return dir
 210             except (OSError, IOError), e:
 211                 if e[0] != _errno.EEXIST:
 212                     break # no point trying more names in this directory
 213                 pass
 214     raise IOError, (_errno.ENOENT,
 215                     ("No usable temporary directory found in %s" % dirlist))
 216
 217 _name_sequence = None
 218
 219 def _get_candidate_names():
 220     """Common setup sequence for all user-callable interfaces."""
 221
 222     global _name_sequence
 223     if _name_sequence is None:
 224         _once_lock.acquire()
 225         try:
 226             if _name_sequence is None:
 227                 _name_sequence = _RandomNameSequence()
 228         finally:
 229             _once_lock.release()
 230     return _name_sequence
 231
 232
 233 def _mkstemp_inner(dir, pre, suf, flags):
 234     """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
 235
 236     names = _get_candidate_names()
 237
 238     for seq in xrange(TMP_MAX):
 239         name = names.next()
 240         file = _os.path.join(dir, pre + name + suf)
 241         try:
 242             fd = _os.open(file, flags, 0600)
 243             _set_cloexec(fd)
 244             return (fd, _os.path.abspath(file))
 245         except OSError, e:
 246             if e.errno == _errno.EEXIST:
 247                 continue # try again
 248             raise
 249
 250     raise IOError, (_errno.EEXIST, "No usable temporary file name found")
 251
 252
 253 # User visible interfaces.
 254
 255 def gettempprefix():
 256     """Accessor for tempdir.template."""
 257     return template
 258
 259 tempdir = None
 260
 261 def gettempdir():
 262     """Accessor for tempdir.tempdir."""
 263     global tempdir
 264     if tempdir is None:
 265         _once_lock.acquire()
 266         try:
 267             if tempdir is None:
 268                 tempdir = _get_default_tempdir()
 269         finally:
 270             _once_lock.release()
 271     return tempdir
 272
 273 def mkstemp(suffix="", prefix=template, dir=None, text=False):
 274     """mkstemp([suffix, [prefix, [dir, [text]]]])
 275     User-callable function to create and return a unique temporary
 276     file.  The return value is a pair (fd, name) where fd is the
 277     file descriptor returned by os.open, and name is the filename.
 278
 279     If 'suffix' is specified, the file name will end with that suffix,
 280     otherwise there will be no suffix.
 281
 282     If 'prefix' is specified, the file name will begin with that prefix,
 283     otherwise a default prefix is used.
 284
 285     If 'dir' is specified, the file will be created in that directory,
 286     otherwise a default directory is used.
 287
 288     If 'text' is specified and true, the file is opened in text
 289     mode.  Else (the default) the file is opened in binary mode.  On
 290     some operating systems, this makes no difference.
 291
 292     The file is readable and writable only by the creating user ID.
 293     If the operating system uses permission bits to indicate whether a
 294     file is executable, the file is executable by no one. The file
 295     descriptor is not inherited by children of this process.
 296
 297     Caller is responsible for deleting the file when done with it.
 298     """
 299
 300     if dir is None:
 301         dir = gettempdir()
 302
 303     if text:
 304         flags = _text_openflags
 305     else:
 306         flags = _bin_openflags
 307
 308     return _mkstemp_inner(dir, prefix, suffix, flags)
 309
 310
 311 def mkdtemp(suffix="", prefix=template, dir=None):
 312     """mkdtemp([suffix, [prefix, [dir]]])
 313     User-callable function to create and return a unique temporary
 314     directory.  The return value is the pathname of the directory.
 315
 316     Arguments are as for mkstemp, except that the 'text' argument is
 317     not accepted.
 318
 319     The directory is readable, writable, and searchable only by the
 320     creating user.
 321
 322     Caller is responsible for deleting the directory when done with it.
 323     """
 324
 325     if dir is None:
 326         dir = gettempdir()
 327
 328     names = _get_candidate_names()
 329
 330     for seq in xrange(TMP_MAX):
 331         name = names.next()
 332         file = _os.path.join(dir, prefix + name + suffix)
 333         try:
 334             _os.mkdir(file, 0700)
 335             return file
 336         except OSError, e:
 337             if e.errno == _errno.EEXIST:
 338                 continue # try again
 339             raise
 340
 341     raise IOError, (_errno.EEXIST, "No usable temporary directory name found")
 342
 343 def mktemp(suffix="", prefix=template, dir=None):
 344     """mktemp([suffix, [prefix, [dir]]])
 345     User-callable function to return a unique temporary file name.  The
 346     file is not created.
 347
 348     Arguments are as for mkstemp, except that the 'text' argument is
 349     not accepted.
 350
 351     This function is unsafe and should not be used.  The file name
 352     refers to a file that did not exist at some point, but by the time
 353     you get around to creating it, someone else may have beaten you to
 354     the punch.
 355     """
 356
 357 ##    from warnings import warn as _warn
 358 ##    _warn("mktemp is a potential security risk to your program",
 359 ##          RuntimeWarning, stacklevel=2)
 360
 361     if dir is None:
 362         dir = gettempdir()
 363
 364     names = _get_candidate_names()
 365     for seq in xrange(TMP_MAX):
 366         name = names.next()
 367         file = _os.path.join(dir, prefix + name + suffix)
 368         if not _exists(file):
 369             return file
 370
 371     raise IOError, (_errno.EEXIST, "No usable temporary filename found")
 372
 373 class _TemporaryFileWrapper:
 374     """Temporary file wrapper
 375
 376     This class provides a wrapper around files opened for
 377     temporary use.  In particular, it seeks to automatically
 378     remove the file when it is no longer needed.
 379     """
 380
 381     def __init__(self, file, name, delete=True):
 382         self.file = file
 383         self.name = name
 384         self.close_called = False
 385         self.delete = delete
 386
 387     def __getattr__(self, name):
 388         file = self.__dict__['file']
 389         a = getattr(file, name)
 390         if type(a) != type(0):
 391             setattr(self, name, a)
 392         return a
 393
 394     # NT provides delete-on-close as a primitive, so we don't need
 395     # the wrapper to do anything special.  We still use it so that
 396     # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
 397     if _os.name != 'nt':
 398
 399         # Cache the unlinker so we don't get spurious errors at
 400         # shutdown when the module-level "os" is None'd out.  Note
 401         # that this must be referenced as self.unlink, because the
 402         # name TemporaryFileWrapper may also get None'd out before
 403         # __del__ is called.
 404         unlink = _os.unlink
 405
 406         def close(self):
 407             if not self.close_called:
 408                 self.close_called = True
 409                 self.file.close()
 410                 if self.delete:
 411                     self.unlink(self.name)
 412
 413         def __del__(self):
 414             self.close()
 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     # file protocol
 515     def __iter__(self):
 516         return self._file.__iter__()
 517
 518     def close(self):
 519         self._file.close()
 520
 521     @property
 522     def closed(self):
 523         return self._file.closed
 524
 525     @property
 526     def encoding(self):
 527         return self._file.encoding
 528
 529     def fileno(self):
 530         self.rollover()
 531         return self._file.fileno()
 532
 533     def flush(self):
 534         self._file.flush()
 535
 536     def isatty(self):
 537         return self._file.isatty()
 538
 539     @property
 540     def mode(self):
 541         return self._file.mode
 542
 543     @property
 544     def name(self):
 545         return self._file.name
 546
 547     @property
 548     def newlines(self):
 549         return self._file.newlines
 550
 551     def next(self):
 552         return self._file.next
 553
 554     def read(self, *args):
 555         return self._file.read(*args)
 556
 557     def readline(self, *args):
 558         return self._file.readline(*args)
 559
 560     def readlines(self, *args):
 561         return self._file.readlines(*args)
 562
 563     def seek(self, *args):
 564         self._file.seek(*args)
 565
 566     @property
 567     def softspace(self):
 568         return self._file.softspace
 569
 570     def tell(self):
 571         return self._file.tell()
 572
 573     def truncate(self):
 574         self._file.truncate()
 575
 576     def write(self, s):
 577         file = self._file
 578         rv = file.write(s)
 579         self._check(file)
 580         return rv
 581
 582     def writelines(self, iterable):
 583         file = self._file
 584         rv = file.writelines(iterable)
 585         self._check(file)
 586         return rv
 587
 588     def xreadlines(self, *args):
 589         return self._file.xreadlines(*args)