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 io as _io
  33 import os as _os
  34 import errno as _errno
  35 from random import Random as _Random
  36
  37 if _os.name == 'mac':
  38     import Carbon.Folder as _Folder
  39     import Carbon.Folders as _Folders
  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 == '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 == 'nt':
 161         dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
 162     else:
 163         dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])
 164
 165     # As a last resort, the current directory.
 166     try:
 167         dirlist.append(_os.getcwd())
 168     except (AttributeError, _os.error):
 169         dirlist.append(_os.curdir)
 170
 171     return dirlist
 172
 173 def _get_default_tempdir():
 174     """Calculate the default directory to use for temporary files.
 175     This routine should be called exactly once.
 176
 177     We determine whether or not a candidate temp dir is usable by
 178     trying to create and write to a file in that directory.  If this
 179     is successful, the test file is deleted.  To prevent denial of
 180     service, the name of the test file must be randomized."""
 181
 182     namer = _RandomNameSequence()
 183     dirlist = _candidate_tempdir_list()
 184     flags = _text_openflags
 185
 186     for dir in dirlist:
 187         if dir != _os.curdir:
 188             dir = _os.path.normcase(_os.path.abspath(dir))
 189         # Try only a few names per directory.
 190         for seq in range(100):
 191             name = next(namer)
 192             filename = _os.path.join(dir, name)
 193             try:
 194                 fd = _os.open(filename, flags, 0o600)
 195                 fp = _io.open(fd, 'wb')
 196                 fp.write(b'blat')
 197                 fp.close()
 198                 _os.unlink(filename)
 199                 del fp, fd
 200                 return dir
 201             except (OSError, IOError) as e:
 202                 if e.args[0] != _errno.EEXIST:
 203                     break # no point trying more names in this directory
 204                 pass
 205     raise IOError(_errno.ENOENT,
 206                   "No usable temporary directory found in %s" % dirlist)
 207
 208 _name_sequence = None
 209
 210 def _get_candidate_names():
 211     """Common setup sequence for all user-callable interfaces."""
 212
 213     global _name_sequence
 214     if _name_sequence is None:
 215         _once_lock.acquire()
 216         try:
 217             if _name_sequence is None:
 218                 _name_sequence = _RandomNameSequence()
 219         finally:
 220             _once_lock.release()
 221     return _name_sequence
 222
 223
 224 def _mkstemp_inner(dir, pre, suf, flags):
 225     """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
 226
 227     names = _get_candidate_names()
 228
 229     for seq in range(TMP_MAX):
 230         name = next(names)
 231         file = _os.path.join(dir, pre + name + suf)
 232         try:
 233             fd = _os.open(file, flags, 0o600)
 234             _set_cloexec(fd)
 235             return (fd, _os.path.abspath(file))
 236         except OSError as e:
 237             if e.errno == _errno.EEXIST:
 238                 continue # try again
 239             raise
 240
 241     raise IOError(_errno.EEXIST, "No usable temporary file name found")
 242
 243
 244 # User visible interfaces.
 245
 246 def gettempprefix():
 247     """Accessor for tempdir.template."""
 248     return template
 249
 250 tempdir = None
 251
 252 def gettempdir():
 253     """Accessor for tempfile.tempdir."""
 254     global tempdir
 255     if tempdir is None:
 256         _once_lock.acquire()
 257         try:
 258             if tempdir is None:
 259                 tempdir = _get_default_tempdir()
 260         finally:
 261             _once_lock.release()
 262     return tempdir
 263
 264 def mkstemp(suffix="", prefix=template, dir=None, text=False):
 265     """User-callable function to create and return a unique temporary
 266     file.  The return value is a pair (fd, name) where fd is the
 267     file descriptor returned by os.open, and name is the filename.
 268
 269     If 'suffix' is specified, the file name will end with that suffix,
 270     otherwise there will be no suffix.
 271
 272     If 'prefix' is specified, the file name will begin with that prefix,
 273     otherwise a default prefix is used.
 274
 275     If 'dir' is specified, the file will be created in that directory,
 276     otherwise a default directory is used.
 277
 278     If 'text' is specified and true, the file is opened in text
 279     mode.  Else (the default) the file is opened in binary mode.  On
 280     some operating systems, this makes no difference.
 281
 282     The file is readable and writable only by the creating user ID.
 283     If the operating system uses permission bits to indicate whether a
 284     file is executable, the file is executable by no one. The file
 285     descriptor is not inherited by children of this process.
 286
 287     Caller is responsible for deleting the file when done with it.
 288     """
 289
 290     if dir is None:
 291         dir = gettempdir()
 292
 293     if text:
 294         flags = _text_openflags
 295     else:
 296         flags = _bin_openflags
 297
 298     return _mkstemp_inner(dir, prefix, suffix, flags)
 299
 300
 301 def mkdtemp(suffix="", prefix=template, dir=None):
 302     """User-callable function to create and return a unique temporary
 303     directory.  The return value is the pathname of the directory.
 304
 305     Arguments are as for mkstemp, except that the 'text' argument is
 306     not accepted.
 307
 308     The directory is readable, writable, and searchable only by the
 309     creating user.
 310
 311     Caller is responsible for deleting the directory when done with it.
 312     """
 313
 314     if dir is None:
 315         dir = gettempdir()
 316
 317     names = _get_candidate_names()
 318
 319     for seq in range(TMP_MAX):
 320         name = next(names)
 321         file = _os.path.join(dir, prefix + name + suffix)
 322         try:
 323             _os.mkdir(file, 0o700)
 324             return file
 325         except OSError as e:
 326             if e.errno == _errno.EEXIST:
 327                 continue # try again
 328             raise
 329
 330     raise IOError(_errno.EEXIST, "No usable temporary directory name found")
 331
 332 def mktemp(suffix="", prefix=template, dir=None):
 333     """User-callable function to return a unique temporary file name.  The
 334     file is not created.
 335
 336     Arguments are as for mkstemp, except that the 'text' argument is
 337     not accepted.
 338
 339     This function is unsafe and should not be used.  The file name
 340     refers to a file that did not exist at some point, but by the time
 341     you get around to creating it, someone else may have beaten you to
 342     the punch.
 343     """
 344
 345 ##    from warnings import warn as _warn
 346 ##    _warn("mktemp is a potential security risk to your program",
 347 ##          RuntimeWarning, stacklevel=2)
 348
 349     if dir is None:
 350         dir = gettempdir()
 351
 352     names = _get_candidate_names()
 353     for seq in range(TMP_MAX):
 354         name = next(names)
 355         file = _os.path.join(dir, prefix + name + suffix)
 356         if not _exists(file):
 357             return file
 358
 359     raise IOError(_errno.EEXIST, "No usable temporary filename found")
 360
 361
 362 class _TemporaryFileWrapper:
 363     """Temporary file wrapper
 364
 365     This class provides a wrapper around files opened for
 366     temporary use.  In particular, it seeks to automatically
 367     remove the file when it is no longer needed.
 368     """
 369
 370     def __init__(self, file, name, delete=True):
 371         self.file = file
 372         self.name = name
 373         self.close_called = False
 374         self.delete = delete
 375
 376     def __getattr__(self, name):
 377         # Attribute lookups are delegated to the underlying file
 378         # and cached for non-numeric results
 379         # (i.e. methods are cached, closed and friends are not)
 380         file = self.__dict__['file']
 381         a = getattr(file, name)
 382         if not isinstance(a, int):
 383             setattr(self, name, a)
 384         return a
 385
 386     # The underlying __enter__ method returns the wrong object
 387     # (self.file) so override it to return the wrapper
 388     def __enter__(self):
 389         self.file.__enter__()
 390         return self
 391
 392     # iter() doesn't use __getattr__ to find the __iter__ method
 393     def __iter__(self):
 394         return iter(self.file)
 395
 396     # NT provides delete-on-close as a primitive, so we don't need
 397     # the wrapper to do anything special.  We still use it so that
 398     # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
 399     if _os.name != 'nt':
 400         # Cache the unlinker so we don't get spurious errors at
 401         # shutdown when the module-level "os" is None'd out.  Note
 402         # that this must be referenced as self.unlink, because the
 403         # name TemporaryFileWrapper may also get None'd out before
 404         # __del__ is called.
 405         unlink = _os.unlink
 406
 407         def close(self):
 408             if not self.close_called:
 409                 self.close_called = True
 410                 self.file.close()
 411                 if self.delete:
 412                     self.unlink(self.name)
 413
 414         def __del__(self):
 415             self.close()
 416
 417         # Need to trap __exit__ as well to ensure the file gets
 418         # deleted when used in a with statement
 419         def __exit__(self, exc, value, tb):
 420             result = self.file.__exit__(exc, value, tb)
 421             self.close()
 422             return result
 423
 424
 425 def NamedTemporaryFile(mode='w+b', buffering=-1, encoding=None,
 426                        newline=None, suffix="", prefix=template,
 427                        dir=None, delete=True):
 428     """Create and return a temporary file.
 429     Arguments:
 430     'prefix', 'suffix', 'dir' -- as for mkstemp.
 431     'mode' -- the mode argument to io.open (default "w+b").
 432     'buffering' -- the buffer size argument to io.open (default -1).
 433     'encoding' -- the encoding argument to io.open (default None)
 434     'newline' -- the newline argument to io.open (default None)
 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 = _io.open(fd, mode, buffering=buffering,
 458                     newline=newline, encoding=encoding)
 459
 460     return _TemporaryFileWrapper(file, name, delete)
 461
 462 if _os.name != 'posix' or _os.sys.platform == 'cygwin':
 463     # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
 464     # while it is open.
 465     TemporaryFile = NamedTemporaryFile
 466
 467 else:
 468     def TemporaryFile(mode='w+b', buffering=-1, encoding=None,
 469                       newline=None, suffix="", prefix=template,
 470                       dir=None):
 471         """Create and return a temporary file.
 472         Arguments:
 473         'prefix', 'suffix', 'dir' -- as for mkstemp.
 474         'mode' -- the mode argument to io.open (default "w+b").
 475         'buffering' -- the buffer size argument to io.open (default -1).
 476         'encoding' -- the encoding argument to io.open (default None)
 477         'newline' -- the newline argument to io.open (default None)
 478         The file is created as mkstemp() would do it.
 479
 480         Returns an object with a file-like interface.  The file has no
 481         name, and will cease to exist when it is closed.
 482         """
 483
 484         if dir is None:
 485             dir = gettempdir()
 486
 487         if 'b' in mode:
 488             flags = _bin_openflags
 489         else:
 490             flags = _text_openflags
 491
 492         (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
 493         try:
 494             _os.unlink(name)
 495             return _io.open(fd, mode, buffering=buffering,
 496                             newline=newline, encoding=encoding)
 497         except:
 498             _os.close(fd)
 499             raise
 500
 501 class SpooledTemporaryFile:
 502     """Temporary file wrapper, specialized to switch from
 503     StringIO to a real file when it exceeds a certain size or
 504     when a fileno is needed.
 505     """
 506     _rolled = False
 507
 508     def __init__(self, max_size=0, mode='w+b', buffering=-1,
 509                  encoding=None, newline=None,
 510                  suffix="", prefix=template, dir=None):
 511         if 'b' in mode:
 512             self._file = _io.BytesIO()
 513         else:
 514             # Setting newline="\n" avoids newline translation;
 515             # this is important because otherwise on Windows we'd
 516             # hget double newline translation upon rollover().
 517             self._file = _io.StringIO(encoding=encoding, newline="\n")
 518         self._max_size = max_size
 519         self._rolled = False
 520         self._TemporaryFileArgs = {'mode': mode, 'buffering': buffering,
 521                                    'suffix': suffix, 'prefix': prefix,
 522                                    'encoding': encoding, 'newline': newline,
 523                                    'dir': dir}
 524
 525     def _check(self, file):
 526         if self._rolled: return
 527         max_size = self._max_size
 528         if max_size and file.tell() > max_size:
 529             self.rollover()
 530
 531     def rollover(self):
 532         if self._rolled: return
 533         file = self._file
 534         newfile = self._file = TemporaryFile(**self._TemporaryFileArgs)
 535         del self._TemporaryFileArgs
 536
 537         newfile.write(file.getvalue())
 538         newfile.seek(file.tell(), 0)
 539
 540         self._rolled = True
 541
 542     # The method caching trick from NamedTemporaryFile
 543     # won't work here, because _file may change from a
 544     # _StringIO instance to a real file. So we list
 545     # all the methods directly.
 546
 547     # Context management protocol
 548     def __enter__(self):
 549         if self._file.closed:
 550             raise ValueError("Cannot enter context with closed file")
 551         return self
 552
 553     def __exit__(self, exc, value, tb):
 554         self._file.close()
 555
 556     # file protocol
 557     def __iter__(self):
 558         return self._file.__iter__()
 559
 560     def close(self):
 561         self._file.close()
 562
 563     @property
 564     def closed(self):
 565         return self._file.closed
 566
 567     @property
 568     def encoding(self):
 569         return self._file.encoding
 570
 571     def fileno(self):
 572         self.rollover()
 573         return self._file.fileno()
 574
 575     def flush(self):
 576         self._file.flush()
 577
 578     def isatty(self):
 579         return self._file.isatty()
 580
 581     @property
 582     def mode(self):
 583         return self._file.mode
 584
 585     @property
 586     def name(self):
 587         return self._file.name
 588
 589     @property
 590     def newlines(self):
 591         return self._file.newlines
 592
 593     def next(self):
 594         return self._file.next
 595
 596     def read(self, *args):
 597         return self._file.read(*args)
 598
 599     def readline(self, *args):
 600         return self._file.readline(*args)
 601
 602     def readlines(self, *args):
 603         return self._file.readlines(*args)
 604
 605     def seek(self, *args):
 606         self._file.seek(*args)
 607
 608     @property
 609     def softspace(self):
 610         return self._file.softspace
 611
 612     def tell(self):
 613         return self._file.tell()
 614
 615     def truncate(self):
 616         self._file.truncate()
 617
 618     def write(self, s):
 619         file = self._file
 620         rv = file.write(s)
 621         self._check(file)
 622         return rv
 623
 624     def writelines(self, iterable):
 625         file = self._file
 626         rv = file.writelines(iterable)
 627         self._check(file)
 628         return rv
 629
 630     def xreadlines(self, *args):
 631         return self._file.xreadlines(*args)